coio/examples/tcp_echo_server-context_pool.cpp at main · Cra3z/coio · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
#include <coio/core.h>
#include <coio/asyncio/io.h>
#include <coio/net/socket.h>
#include <coio/net/tcp.h>
#include <coio/utils/signal_wait.h>
#include "common.h"

#if COIO_OS_LINUX
#include <coio/asyncio/epoll_context.h>
using io_context = coio::epoll_context;
#elif COIO_OS_WINDOWS
#include <coio/asyncio/iocp_context.h>
using io_context = coio::iocp_context;
#endif

using tcp_socket = coio::tcp::socket<io_context::scheduler>;
using tcp_acceptor = coio::tcp::acceptor<io_context::scheduler>;

class io_context_pool {
public:
    explicit io_context_pool(std::size_t count) {
        COIO_ASSERT(count > 0);
        io_contexts_.reserve(count);
        work_guards_.reserve(count);
        threads_.reserve(count);
        for (std::size_t i = 0; i < count; ++i) {
            work_guards_.emplace_back(*io_contexts_.emplace_back(std::make_unique<io_context>()));
        }

        for (std::size_t i = 0; i < count; ++i) {
            threads_.emplace_back([this, i] {
                ::debug("worker started");
                io_contexts_[i]->run();
                ::debug("worker finished");
            });
        }
    }

    io_context_pool(const io_context_pool&) = delete;

    io_context_pool& operator= (const io_context_pool&) = delete;

    ~io_context_pool() {
        stop();
        for (auto& thread : threads_) {
            if (thread.joinable()) {
                thread.join();
            }
        }
    }

    auto stop() -> void {
        for (auto& ctx : io_contexts_) {
            ctx->request_stop();
        }
        work_guards_.clear();
    }

    auto pick_scheduler() noexcept -> io_context::scheduler {
        return io_contexts_[std::exchange(next_, (next_ + 1) % io_contexts_.size())]->get_scheduler();
    }

private:
    std::size_t next_ = 0;
    std::vector<std::unique_ptr<io_context>> io_contexts_;
    std::vector<std::thread> threads_;
    std::vector<coio::work_guard<io_context>> work_guards_;
};

auto handle_connection(tcp_socket socket) -> io_context::task<> {
    auto remote_endpoint = socket.remote_endpoint();
    ::debug("new connection from [{}]", remote_endpoint);
    try {
        char buffer[1024];
        while (true) {
            const auto length = co_await socket.async_read_some(coio::as_writable_bytes(buffer));
            co_await (coio::async_write(socket, coio::as_bytes(buffer, length)) | as_throwing);
        }
    }
    catch (const std::system_error& e) {
        ::debug("connection with [{}] broken because \"{}\"", remote_endpoint, e.what());
    }
}

auto start_server(io_context_pool& pool, coio::async_scope& scope) -> io_context::task<> try {
    tcp_acceptor acceptor{co_await coio::read_scheduler(), coio::endpoint{coio::ipv4_address::any(), 8086}};
    ::debug("server \"{}\" start...", acceptor.local_endpoint());
    while (true) {
        auto next_scheduler = pool.pick_scheduler();
        scope.spawn_on(next_scheduler, handle_connection(co_await acceptor.async_accept(next_scheduler)));
    }
}
catch (const std::system_error& e) {
    ::debug("acceptor error: {}", e.what());
}

auto signal_watchdog(io_context_pool& pool) -> coio::inline_task<> {
    const int signum = co_await coio::signal_wait(SIGINT, SIGTERM);
    ::debug("server stop with signal: ({}){}", signum, coio::strsignal(signum));
    pool.stop();
}

auto main() -> int {
    io_context_pool pool{4};
    coio::async_scope scope;
    scope.spawn(signal_watchdog(pool));
    scope.spawn_on(pool.pick_scheduler(), start_server(pool, scope));
    coio::this_thread::sync_wait(scope.join());
}