proactor_pool.h

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

#pragma once

#include <pthread.h>

#include <vector>

#include "absl/container/flat_hash_set.h"
#include "absl/strings/string_view.h"
#include "base/arena.h"
#include "base/type_traits.h"
#include "base/RWSpinLock.h"
#include "util/uring/proactor.h"

namespace util {
namespace uring {

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

 public:
  ProactorPool(const ProactorPool&) = delete;
  void operator=(const ProactorPool&) = delete;

  explicit ProactorPool(std::size_t pool_size = 0);

  ~ProactorPool();

  void Run(uint32_t ring_depth = 256);

  void Stop();

  Proactor* GetNextProactor();

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

  Proactor& at(size_t i) {
    return proactor_[i];
  }

  size_t size() const {
    return pool_size_;
  }

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

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

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

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

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

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

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

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

  Proactor* GetLocalProactor();

  // Auxillary functions

  // Returns a string owned by pool's global storage. Allocates only once for each new string blob.
  // Currently has average performance and it employs RW spinlock underneath.
  absl::string_view GetString(absl::string_view source);

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

  std::unique_ptr<Proactor[]> proactor_;

  std::atomic_uint_fast32_t next_io_context_{0};
  uint32_t pool_size_;

  folly::RWSpinLock str_lock_;
  absl::flat_hash_set<absl::string_view> str_set_;
  base::Arena arena_;

  enum State { STOPPED, RUN } state_ = STOPPED;
};

constexpr size_t foo = sizeof(ProactorPool);

}  // namespace uring
}  // namespace util