heavyworker.cc

#include <chrono>
#include <iostream>
#include <memory>
#include <thread>
#include <vector>

#include <boost/asio.hpp>

using namespace std::chrono_literals;

class server;

class session : public std::enable_shared_from_this<session> {
public:
  session(std::weak_ptr<server> server, boost::asio::ip::tcp::socket socket)
      : server_{std::move(server)}, socket_{std::move(socket)} {}

  void start() {
    do_read();
  }

private:
  void do_reply() {
    auto self = shared_from_this();
    socket_.async_write_some(
        boost::asio::buffer("thanks!\n"),
        [this, self](boost::system::error_code ec, [[maybe_unused]] std::size_t length) {
          if (!ec) {
            do_read();
          }
        });
  }

  void do_read();

  std::weak_ptr<server> server_;
  boost::asio::ip::tcp::socket socket_;

  char readbuf_[0x10];
};

class server : public std::enable_shared_from_this<server> {
public:
  server(boost::asio::io_context& listen_ctx, boost::asio::io_context& worker_ctx)
      : listen_ctx_{listen_ctx},
        worker_ctx_{worker_ctx},
        work_{worker_ctx_},
        timer_{listen_ctx_},
        acceptor_{listen_ctx_, boost::asio::ip::tcp::endpoint{boost::asio::ip::tcp::v4(), 4567}} {
    do_accept();

    timer_.expires_after(10s);
    do_timer();
  }

private:
  void do_accept() {
    acceptor_.async_accept(
        [this](boost::system::error_code ec, boost::asio::ip::tcp::socket socket) {
          if (!ec) {
            std::make_shared<session>(shared_from_this(), std::move(socket))->start();
          }
          do_accept();
        });
  }

  void do_timer() {
    timer_.async_wait([this](boost::system::error_code ec) {
      if (ec) return;

      swap_buffer(); // buffer の入れ替えは listen_ctx_ のスレッドで実行

      boost::asio::post(worker_ctx_, [this, w = boost::asio::io_context::work(listen_ctx_)] {
        timer_.expires_after(10s);
        do_work(buffer_back());
        do_timer();
      });
    });
  }

  void do_work(std::vector<int>& buffer) {
    std::cout << "worker: start heavy work @ thread " << std::this_thread::get_id() << '\n';

    std::this_thread::sleep_for(5s);

    for (const auto& v : buffer) {
      std::cout << "worker: " << v << '\n';
    }
    buffer.clear();

    std::cout << "worker: done!\n";
  }

  std::vector<int>& buffer_front() {
    return buffer_[buffer_state_];
  }

  std::vector<int>& buffer_back() {
    return buffer_[!buffer_state_];
  }

  void swap_buffer() {
    buffer_state_ = !buffer_state_;
  }

  boost::asio::io_context& listen_ctx_;
  boost::asio::io_context& worker_ctx_;

  boost::asio::io_context::work work_;
  boost::asio::steady_timer timer_;
  boost::asio::ip::tcp::acceptor acceptor_;

  bool buffer_state_ = false;
  std::vector<int> buffer_[2];

  friend class session;
};

void session::do_read() {
  auto self = shared_from_this();
  socket_.async_read_some(boost::asio::buffer(readbuf_, sizeof readbuf_ - 1),
                          [this, self](boost::system::error_code ec, std::size_t length) {
                            if (ec) return;

                            if (auto server = server_.lock(); server) {
                              readbuf_[length] = '\0';
                              server->buffer_front().push_back(std::atoi(readbuf_));
                              do_reply();
                            }
                          });
}

auto main() -> int {
  boost::asio::io_context listen_ctx{1};
  boost::asio::io_context worker_ctx{1};

  std::cout << "main thread: " << std::this_thread::get_id() << std::endl;
  boost::asio::post(worker_ctx, [] {
    std::cout << "worker thread: " << std::this_thread::get_id() << std::endl;
  });

  auto s = std::make_shared<server>(listen_ctx, worker_ctx);

  std::thread worker{
      static_cast<std::size_t (boost::asio::io_context::*)()>(&boost::asio::io_context::run),
      &worker_ctx};

  boost::asio::signal_set signal(listen_ctx, SIGINT, SIGTERM);
  signal.async_wait([&listen_ctx, &worker_ctx, &worker](boost::system::error_code ec,
                                                        [[maybe_unused]] int signal) {
    if (!ec) {
      std::cout << "\nstopping..." << std::endl;

      worker_ctx.stop();
      worker.join();

      listen_ctx.stop();
    }
  });

  listen_ctx.run();
}