io_context_pool.h

// Copyright 2018, Beeri 15.  All rights reserved.
// Author: Roman Gershman (romange@gmail.com)
//

#pragma once

#include <thread>
#include <vector>

#include <absl/types/optional.h>
#include <boost/asio/executor_work_guard.hpp>
#include <boost/fiber/fiber.hpp>

#include "base/type_traits.h"
#include "util/asio/io_context.h"
#include "util/fibers/fibers_ext.h"

namespace util {

class IoContextPool {
  template <typename Func, typename... Args>
  using AcceptArgsCheck =
      typename std::enable_if<base::is_invocable<Func, Args...>::value, int>::type;

 public:
  using io_context = ::boost::asio::io_context;

  IoContextPool(const IoContextPool&) = delete;
  void operator=(const IoContextPool&) = delete;

  explicit IoContextPool(std::size_t pool_size = 0, std::vector<size_t> cpus = {});

  ~IoContextPool();

  void Run();

  void Stop();

  IoContext& GetNextContext();

  IoContext& operator[](size_t i) { return context_arr_[i]; }
  IoContext& at(size_t i) { return context_arr_[i]; }

  size_t size() const { return context_arr_.size(); }

  template <typename Func, AcceptArgsCheck<Func, IoContext&> = 0> void AsyncOnAll(Func&& func) {
    CheckRunningState();
    for (unsigned i = 0; i < size(); ++i) {
      IoContext& context = context_arr_[i];
      // func must be copied, it can not be moved, because we dsitribute it into multiple
      // IoContexts.
      context.Async([&context, func] () mutable { func(context); });
    }
  }

  template <typename Func, AcceptArgsCheck<Func, unsigned, IoContext&> = 0>
  void AsyncOnAll(Func&& func) {
    CheckRunningState();
    for (unsigned i = 0; i < size(); ++i) {
      IoContext& context = context_arr_[i];
      // Copy func on purpose, see above.
      context.Async([&context, i, func] () mutable { func(i, context); });
    }
  }

  template <typename Func, AcceptArgsCheck<Func, IoContext&> = 0> void AwaitOnAll(Func&& func) {
    fibers_ext::BlockingCounter bc(size());
    auto cb = [func = std::forward<Func>(func), bc](IoContext& context) mutable {
      func(context);
      bc.Dec();
    };
    AsyncOnAll(std::move(cb));
    bc.Wait();
  }

  template <typename Func, AcceptArgsCheck<Func, unsigned, IoContext&> = 0>
  void AwaitOnAll(Func&& func) {
    fibers_ext::BlockingCounter bc(size());
    auto cb = [func = std::forward<Func>(func), bc](unsigned index, IoContext& context) mutable {
      func(index, context);
      bc.Dec();
    };
    AsyncOnAll(std::move(cb));
    bc.Wait();
  }

  template <typename Func, AcceptArgsCheck<Func, unsigned, IoContext&> = 0> void AsyncFiberOnAll(Func&& func) {
    AsyncOnAll([func = std::forward<Func>(func)](unsigned i, IoContext& context) {
      ::boost::fibers::fiber(func, i, std::ref(context)).detach();
    });
  }

  template <typename Func, AcceptArgsCheck<Func, IoContext&> = 0> void AsyncFiberOnAll(Func&& func) {
    AsyncOnAll([func = std::forward<Func>(func)](IoContext& context) {
        ::boost::fibers::fiber(func, std::ref(context)).detach();
    });
  }

  template <typename Func, AcceptArgsCheck<Func, unsigned, IoContext&> = 0>
  void AwaitFiberOnAll(Func&& func) {
    fibers_ext::BlockingCounter bc(size());
    auto cb = [func = std::forward<Func>(func), bc](unsigned i, IoContext& context) mutable {
      func(i, context);
      bc.Dec();
    };
    AsyncFiberOnAll(std::move(cb));
    bc.Wait();
  }

  template <typename Func, AcceptArgsCheck<Func, IoContext&> = 0> void AwaitFiberOnAll(Func&& func) {
    fibers_ext::BlockingCounter bc(size());
    auto cb = [func = std::forward<Func>(func), bc](IoContext& context) mutable {
      func(context);
      bc.Dec();
    };
    AsyncFiberOnAll(std::move(cb));
    bc.Wait();
  }

  template <typename Func> void AwaitFiberOnAllSerially(Func&& func) {
    boost::fibers::mutex mu; // Enforces only one thread.
    AwaitFiberOnAll([&mu,func](auto&& args) { // We copy func on purpose, see ebove.
        std::lock_guard<boost::fibers::mutex> lk(mu);
        func(std::forward<decltype(args)>(args));
    });
  }

  IoContext* GetThisContext();

 private:
  void WrapLoop(size_t index, fibers_ext::BlockingCounter* bc);
  void CheckRunningState();

  typedef ::boost::asio::executor_work_guard<IoContext::io_context::executor_type> work_guard_t;

  std::vector<size_t> cpu_idx_arr_;
  std::vector<IoContext> context_arr_;
  struct TInfo {
    pthread_t tid = 0;
    absl::optional<work_guard_t> work;
  };

  std::vector<TInfo> thread_arr_;

  std::atomic_uint_fast32_t next_io_context_{0};
  thread_local static size_t context_indx_;
  enum State { STOPPED, RUN } state_ = STOPPED;
};

}  // namespace util