io-service.hpp

#pragma once

#include <coroutine>
#include <iostream>
#include <utility>
#include <queue>
#include <unordered_map>
#include <format>
#include <functional>
#include <list>
#include <sys/time.h>
#include <sys/timerfd.h>
#include <sys/epoll.h>
#include <sys/eventfd.h>
#include <sys/socket.h>
#include <sys/timerfd.h>
#include <sys/types.h>
#include "task.hpp"
#include "unistd.h"

class io_service
{
public:
    using task_t = task<int>;
    using handle_t = std::coroutine_handle<typename task_t::promise_type>;
    using schedule_t = std::function<void(handle_t)>;

    enum class poll_op : uint64_t
    {
        /// Poll for read operations.
        read = EPOLLIN,
        /// Poll for write operations.
        write = EPOLLOUT,
        /// Poll for read and write operations.
        read_write = EPOLLIN | EPOLLOUT
    };

    struct poll_info
    {
        int fd;
        handle_t h;
    };

    explicit io_service(schedule_t schedule)
        : m_schedule(schedule),
          m_epoll_fd(epoll_create1(EPOLL_CLOEXEC)),
          m_timer_fd(timerfd_create(CLOCK_MONOTONIC, TFD_NONBLOCK | TFD_CLOEXEC))
    {
        epoll_event e{};
        e.data.ptr = const_cast<void *>(m_timer_ptr);
        epoll_ctl(m_epoll_fd, EPOLL_CTL_ADD, m_timer_fd, &e);
        std::cout << "m_epoll_fd = " << m_epoll_fd << std::endl;
        std::cout << "m_timer_fd = " << m_timer_fd << std::endl;
    }
    io_service(const io_service &) = delete;
    io_service(io_service &&) = delete;
    auto operator=(const io_service &) -> io_service & = delete;
    auto operator=(io_service &&) -> io_service & = delete;

    struct fd_awaiter
    {
        bool await_ready() { return false; }
        void await_suspend(handle_t h)
        {
            h = get_root_coro(h);
            epoll_event e{};
            e.events = static_cast<uint32_t>(op) | EPOLLONESHOT | EPOLLRDHUP;
            poll_info *pi = new poll_info{};
            pi->fd = fd;
            pi->h = h;
            e.data.ptr = pi;
            //std::cout << "register fd = " << fd << " for " << static_cast<uint32_t>(op) << std::endl;
            if (epoll_ctl(m_epoll_fd, EPOLL_CTL_ADD, fd, &e))
            {
                std::cerr << "epoll ctl error on fd " << fd << "\n";
                if (errno == EEXIST) {
                    std::cerr << "err = eexist" << std::endl;
                }
            }
            coro = h;
            original_state = h.promise().last_await_state;
            h.promise().last_await_state = await_state::no_schedule;
        }
        void await_resume()
        {
            coro.promise().last_await_state = original_state;
        }
        int m_epoll_fd;
        int fd;
        poll_op op;
        handle_t coro;
        await_state original_state;
    };

    auto wait_stdin()
    {
        return fd_awaiter{m_epoll_fd, STDIN_FILENO, poll_op::read};
    }

    auto wait_read(int fd)
    {
        // std::cout << "wait read " << fd << std::endl;
        return fd_awaiter{m_epoll_fd, fd, poll_op::read};
    }

    auto wait_write(int fd)
    {
        return fd_awaiter{m_epoll_fd, fd, poll_op::write};
    }

    auto wait_read_write(int fd)
    {
        return fd_awaiter{m_epoll_fd, fd, poll_op::read_write};
    }
    task_t co_check_io()
    {
        co_yield "co_check_io";
        co_await std::suspend_always{}; // stop at beginning
        while (1)
        {
            // std::cout << "before epoll wait" << std::endl;
            auto event_count = epoll_wait(m_epoll_fd, m_events.data(), m_max_events, -1);
            if (event_count > 0)
            {
                // std::cout << "event count = " << event_count << std::endl;
                for (std::size_t i = 0; i < static_cast<std::size_t>(event_count); ++i)
                {
                    epoll_event &event = m_events[i];
                    void *data_ptr = event.data.ptr;

                    if (data_ptr == m_timer_ptr)
                    {
                        // Process all events that have timed out.
                    }
                    else if (data_ptr != nullptr)
                    {
                        auto pi = static_cast<poll_info *>(data_ptr);
                        // std::cout << "process data_ptr " << pi << std::endl;
                        int fd = pi->fd;
                        handle_t h = pi->h;
                        delete pi;
                        epoll_ctl(m_epoll_fd, EPOLL_CTL_DEL, fd, nullptr);
                        //std::cout << "unregister fd = "  << fd << std::endl;
                        m_schedule(h);
                    }
                }
            }

            co_await std::suspend_always{}; // yield to let other corotines to run
        }
    }

private:
    int m_epoll_fd;
    int m_timer_fd;
    static constexpr const int m_timer_object{0};
    static constexpr const void *m_timer_ptr = &m_timer_object;
    static const constexpr std::size_t m_max_events = 16;
    std::array<struct epoll_event, m_max_events> m_events{};
    schedule_t m_schedule;
};