YugabyteDB (2.13.1.0-b60, 21121d69985fbf76aa6958d8f04a9bfa936293b5)

// Copyright (c) YugaByte, Inc.

//

// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except

// in compliance with the License.  You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software distributed under the License

// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express

// or implied.  See the License for the specific language governing permissions and limitations

// under the License.

//

#ifndef YB_UTIL_TASKSTREAM_H

#define YB_UTIL_TASKSTREAM_H

#include <atomic>

#include <chrono>

#include <condition_variable>

#include <memory>

#include <vector>

#include <gflags/gflags.h>

#include "yb/util/status_fwd.h"

#include "yb/util/blocking_queue.h"

#include "yb/util/status_format.h"

#include "yb/util/thread.h"

#include "yb/util/threadpool.h"

using namespace std::chrono_literals;

using std::vector;

namespace yb {

class ThreadPool;

class ThreadPoolToken;

template <typename T>

class TaskStreamImpl;

YB_DEFINE_ENUM(TaskStreamRunState, (kIdle)(kSubmit)(kDrain)(kProcess)(kComplete)(kFinish));

template <typename T>

// TaskStream has a thread pool token in the given thread pool.

// TaskStream does not manage a thread but only submits to the token in the thread pool.

// When we submit tasks to the taskstream, it adds tasks to the queue to be processed.

// The internal Run function will call the user-provided function,

// for each element in the queue in a loop.

// When the queue is empty, it calls the user-provided function with no parameter,

// to indicate it needs to process the end of the group of tasks processed.

// This feature is used for the preparer and appender functionality.

class TaskStream {

 public:

  explicit TaskStream(std::function<void(T*)> process_item,

                      ThreadPool* thread_pool,

                      int32_t queue_max_size,

                      const MonoDelta& queue_max_wait);

  ~TaskStream();

  CHECKED_STATUS Start();

  void Stop();

  CHECKED_STATUS Submit(T* item);

  CHECKED_STATUS TEST_SubmitFunc(const std::function<void()>& func);

  std::string GetRunThreadStack() {

    auto result = ThreadStack(run_tid_);

    if (!result.ok()) {

      return result.status().ToString();

    }

    return (*result).Symbolize();

  }

  std::string ToString() const {

    return YB_CLASS_TO_STRING(queue, run_state, stopped, stop_requested);

  }

 private:

  using RunState = TaskStreamRunState;

  void ChangeRunState(RunState expected_old_state, RunState new_state) {

    auto old_state = run_state_.exchange(new_state, std::memory_order_acq_rel);

    LOG_IF(DFATAL, old_state != expected_old_state)

        << "Task stream was in wrong state " << old_state << " while "

        << expected_old_state << " was expected";

  }

  // We set this to true to tell the Run function to return. No new tasks will be accepted, but

  // existing tasks will still be processed.

  std::atomic<bool> stop_requested_{false};

  std::atomic<RunState> run_state_{RunState::kIdle};

  // This is set to true immediately before the thread exits.

  std::atomic<bool> stopped_{false};

  // The objects in the queue are owned by the queue and ownership gets tranferred to ProcessItem.

  BlockingQueue<T*> queue_;

  // This mutex/condition combination is used in Stop() in case multiple threads are calling that

  // function concurrently. One of them will ask the taskstream thread to stop and wait for it, and

  // then will notify other threads that have called Stop().

  std::mutex stop_mtx_;

  std::condition_variable stop_cond_;

  std::unique_ptr<ThreadPoolToken> taskstream_pool_token_;

  std::function<void(T*)> process_item_;

  ThreadIdForStack run_tid_ = 0;

  // Maximum time to wait for the queue to become non-empty.

  const MonoDelta queue_max_wait_;

  void Run();

  void ProcessItem(T* item);

};

template <typename T>

TaskStream<T>::TaskStream(std::function<void(T*)> process_item,

                          ThreadPool* thread_pool,

                          int32_t queue_max_size,

                          const MonoDelta& queue_max_wait)

    : queue_(queue_max_size),

      // run_state_ is responsible for serial execution, so we use concurrent here to avoid

      // unnecessary checks in thread pool.

      taskstream_pool_token_(thread_pool->NewToken(ThreadPool::ExecutionMode::CONCURRENT)),

      process_item_(process_item),

      queue_max_wait_(queue_max_wait) {

}

template <typename T>

TaskStream<T>::~TaskStream() {

  Stop();

}

template <typename T>

Status TaskStream<T>::Start() {

  VLOG(1) << "Starting the TaskStream";

  return Status::OK();

}

template <typename T>

void TaskStream<T>::Stop() {

  VLOG(1) << "Stopping the TaskStream";

  auto scope_exit = ScopeExit([] {

    VLOG
(1) << "The TaskStream has stopped"1
;

  });

  queue_.Shutdown();

  if (stopped_.load(std::memory_order_acquire)) {

    return;

  }

  stop_requested_ = true;

  {

    std::unique_lock<std::mutex> stop_lock(stop_mtx_);

    stop_cond_.wait(stop_lock, [this] {

      return (run_state_.load(std::memory_order_acquire) == RunState::kIdle && 
queue_.empty()76.4k
);

    });

  }

  stopped_.store(true, std::memory_order_release);

}

template <typename T>

Status TaskStream<T>::Submit(T *task) {

  if (stop_requested_.load(std::memory_order_acquire)) {

    return STATUS(IllegalState, "Tablet is shutting down");

  }

  if (!queue_.BlockingPut(task)) {

    return STATUS_FORMAT(ServiceUnavailable,

                         "TaskStream queue is full (max capacity $0)",

                         queue_.max_size());

  }

  RunState expected = RunState::kIdle;

  if (!run_state_.compare_exchange_strong(

          expected, RunState::kSubmit, std::memory_order_acq_rel)) {

    // run_state_ was not idle, so we are not creating a task to process operations.

    return Status::OK();

  }

  return taskstream_pool_token_->SubmitFunc(std::bind(&TaskStream::Run, this));

}

template <typename T>

Status TaskStream<T>::TEST_SubmitFunc(const std::function<void()>& func) {

  return taskstream_pool_token_->SubmitFunc(func);

}

template <typename T>

void TaskStream<T>::Run() {

  VLOG
(1) << "Starting taskstream task:" << this6
;

  ChangeRunState(RunState::kSubmit, RunState::kDrain);

  run_tid_ = Thread::CurrentThreadIdForStack();

  for (;;) {

    MonoTime wait_timeout_deadline = MonoTime::Now() + queue_max_wait_;

    std::vector<T *> group;

    queue_.BlockingDrainTo(&group, wait_timeout_deadline);

    if (!group.empty()) {

      ChangeRunState(RunState::kDrain, RunState::kProcess);

      for (T* item : group) {

        ProcessItem(item);

      }

      ChangeRunState(RunState::kProcess, RunState::kComplete);

      ProcessItem(nullptr);

      group.clear();

      ChangeRunState(RunState::kComplete, RunState::kDrain);

      continue;

    }

    ChangeRunState(RunState::kDrain, RunState::kFinish);

    // Not processing and queue empty, return from task.

    std::unique_lock<std::mutex> stop_lock(stop_mtx_);

    ChangeRunState(RunState::kFinish, RunState::kIdle);

    if (!queue_.empty()) {

      // Got more operations, try stay in the loop.

      RunState expected = RunState::kIdle;

      if (run_state_.compare_exchange_strong(

              expected, RunState::kDrain, std::memory_order_acq_rel)) {

        continue;

      }

    }

    if (stop_requested_.load(std::memory_order_acquire)) {

      VLOG(1) << "TaskStream task's Run() function is returning because stop is requested.";

      stop_cond_.notify_all();

      return;

    }

    VLOG
(1) << "Returning from TaskStream task after inactivity:" << this21.2k
;

    return;

  }

}

template <typename T>

void TaskStream<T>::ProcessItem(T* item) {

  process_item_(item);

}

}  // namespace yb

#endif  // YB_UTIL_TASKSTREAM_H