C++ Task::GetIoWaitData方法代码示例

void IoWait::CoSwitch(std::vector<FdStruct> && fdsts, int timeout_ms)
{
    Task* tk = g_Scheduler.GetCurrentTask();
    if (!tk) return ;

    uint32_t id = ++tk->GetIoWaitData().io_block_id_;
    tk->state_ = TaskState::io_block;
    tk->GetIoWaitData().wait_successful_ = 0;
    tk->GetIoWaitData().io_block_timeout_ = timeout_ms;
    tk->GetIoWaitData().io_block_timer_.reset();
    tk->GetIoWaitData().wait_fds_.swap(fdsts);
    for (auto &fdst : tk->GetIoWaitData().wait_fds_) {
        fdst.epoll_ptr.tk = tk;
        fdst.epoll_ptr.io_block_id = id;
    }

    DebugPrint(dbg_ioblock, "task(%s) CoSwitch id=%d, nfds=%d, timeout=%d",
            tk->DebugInfo(), id, (int)fdsts.size(), timeout_ms);
    g_Scheduler.CoYield();
}

int select(int nfds, fd_set *readfds, fd_set *writefds,
        fd_set *exceptfds, struct timeval *timeout)
{
    if (!select_f) coroutine_hook_init();

    int timeout_ms = -1;
    if (timeout)
        timeout_ms = timeout->tv_sec * 1000 + timeout->tv_usec / 1000;

    Task* tk = g_Scheduler.GetCurrentTask();
    DebugPrint(dbg_hook, "task(%s) hook select(nfds=%d, rd_set=%p, wr_set=%p, er_set=%p, timeout=%d ms).",
            tk ? tk->DebugInfo() : "nil",
            (int)nfds, readfds, writefds, exceptfds, timeout_ms);

    if (!tk)
        return select_f(nfds, readfds, writefds, exceptfds, timeout);

    if (timeout_ms == 0)
        return select_f(nfds, readfds, writefds, exceptfds, timeout);

    if (!nfds && !readfds && !writefds && !exceptfds && timeout) {
        g_Scheduler.SleepSwitch(timeout_ms);
        return 0;
    }

    nfds = std::min<int>(nfds, FD_SETSIZE);
    std::pair<fd_set*, uint32_t> sets[3] =
    {
        {readfds, EPOLLIN | EPOLLERR | EPOLLHUP},
        {writefds, EPOLLOUT},
        {exceptfds, EPOLLERR | EPOLLHUP}
    };

    static const char* set_names[] = {"readfds", "writefds", "exceptfds"};
    std::vector<FdStruct> fdsts;
    for (int i = 0; i < nfds; ++i) {
        FdStruct *fdst = NULL;
        for (int si = 0; si < 3; ++si) {
            if (!sets[si].first)
                continue;

            if (!FD_ISSET(i, sets[si].first))
                continue;

            if (!fdst) {
                fdsts.emplace_back();
                fdst = &fdsts.back();
                fdst->fd = i;
            }

            fdsts.back().event |= sets[si].second;
            DebugPrint(dbg_hook, "task(%s) hook select %s(%d)",
                    tk->DebugInfo(), set_names[si], (int)i);
        }
    }

    g_Scheduler.IOBlockSwitch(std::move(fdsts), timeout_ms);
    bool is_timeout = false;
    if (tk->GetIoWaitData().io_block_timer_) {
        is_timeout = true;
        if (g_Scheduler.BlockCancelTimer(tk->GetIoWaitData().io_block_timer_)) {
            is_timeout = false;
            tk->DecrementRef(); // timer use ref.
        }
    }

    if (tk->GetIoWaitData().wait_successful_ == 0) {
        if (is_timeout) {
            if (readfds) FD_ZERO(readfds);
            if (writefds) FD_ZERO(writefds);
            if (exceptfds) FD_ZERO(exceptfds);
            return 0;
        } else {
            if (timeout_ms > 0)
                g_Scheduler.SleepSwitch(timeout_ms);
            timeval immedaitely = {0, 0};
            return select_f(nfds, readfds, writefds, exceptfds, &immedaitely);
        }
    }

    int n = 0;
    for (auto &fdst : tk->GetIoWaitData().wait_fds_) {
        int fd = fdst.fd;
        for (int si = 0; si < 3; ++si) {
            if (!sets[si].first)
                continue;

            if (!FD_ISSET(fd, sets[si].first))
                continue;

            if (sets[si].second & fdst.epoll_ptr.revent) {
                ++n;
                continue;
            }

            FD_CLR(fd, sets[si].first);
        }
    }

    return n;
//.........这里部分代码省略.........

int poll(struct pollfd *fds, nfds_t nfds, int timeout)
{
    if (!poll_f) coroutine_hook_init();

    Task* tk = g_Scheduler.GetCurrentTask();
    DebugPrint(dbg_hook, "task(%s) hook poll(nfds=%d, timeout=%d). %s coroutine.",
            tk ? tk->DebugInfo() : "nil",
            (int)nfds, timeout,
            g_Scheduler.IsCoroutine() ? "In" : "Not in");

    if (!g_Scheduler.IsCoroutine())
        return poll_f(fds, nfds, timeout);

    if (timeout == 0)
        return poll_f(fds, nfds, timeout);

    if (nfds == 0) {
        // co sleep
        g_Scheduler.SleepSwitch(timeout);
        return 0;
    }

    std::vector<FdStruct> fdsts;
    for (nfds_t i = 0; i < nfds; ++i) {
        fdsts.emplace_back();
        fdsts.back().fd = fds[i].fd;
        fdsts.back().event = PollEvent2Epoll(fds[i].events);
        DebugPrint(dbg_hook, "hook poll task(%s), fd[%d]=%d.",
                tk->DebugInfo(), (int)i, fds[i].fd);
    }

    // add into epoll, and switch other context.
    g_Scheduler.IOBlockSwitch(std::move(fdsts), timeout);
    bool is_timeout = false; // 是否超时
    if (tk->GetIoWaitData().io_block_timer_) {
        is_timeout = true;
        if (g_Scheduler.BlockCancelTimer(tk->GetIoWaitData().io_block_timer_)) {
            tk->DecrementRef(); // timer use ref.
            is_timeout = false;
        }
    }

    if (tk->GetIoWaitData().wait_successful_ == 0) {
        if (is_timeout)
            return 0;
        else {
            // 加入epoll失败
            if (timeout > 0)
                g_Scheduler.SleepSwitch(timeout);
            return poll_f(fds, nfds, 0);
        }
    }

    int n = 0;
    for (int i = 0; i < (int)tk->GetIoWaitData().wait_fds_.size(); ++i)
    {
        fds[i].revents = EpollEvent2Poll(tk->GetIoWaitData().wait_fds_[i].epoll_ptr.revent);
        if (fds[i].revents) ++n;
    }

//    /// 在一次epoll_wait调用中, 同一个fd可能会被触发多次, 此处不必做严格校验
//    if (n != (int)tk->GetIoWaitData().wait_successful_)
//    {
//        DebugPrint(dbg_debugger, "task(%s) poll assert. n=%d, "
//                "wait_successful=%d, fds_size=%d",
//                tk->DebugInfo(),
//                n, (int)tk->GetIoWaitData().wait_successful_,
//                (int)tk->GetIoWaitData().wait_fds_.size());
//
//        for (int i = 0; i < (int)tk->GetIoWaitData().wait_fds_.size(); ++i)
//        {
//            fds[i].revents = EpollEvent2Poll(tk->GetIoWaitData().wait_fds_[i].epoll_ptr.revent);
//            DebugPrint(dbg_debugger, "[%d] epoll_event=%d, poll_event=%d",
//                    i, tk->GetIoWaitData().wait_fds_[i].epoll_ptr.revent, fds[i].revents);
//        }
//    }
//    assert(n == (int)tk->GetIoWaitData().wait_successful_);

    return n;
}

static ssize_t read_write_mode(int fd, OriginF fn, const char* hook_fn_name, uint32_t event, int timeout_so, Args && ... args)
{
    Task* tk = g_Scheduler.GetCurrentTask();
    DebugPrint(dbg_hook, "task(%s) hook %s. %s coroutine.",
            tk ? tk->DebugInfo() : "nil", hook_fn_name, g_Scheduler.IsCoroutine() ? "In" : "Not in");

    if (!tk)
        return fn(fd, std::forward<Args>(args)...);

    struct stat fd_stat;
    if (-1 == fstat(fd, &fd_stat))
        return fn(fd, std::forward<Args>(args)...);

    if (!S_ISSOCK(fd_stat.st_mode)) // 不是socket, 不HOOK.
        return fn(fd, std::forward<Args>(args)...);

    int flags = fcntl(fd, F_GETFL, 0);
    if (-1 == flags || (flags & O_NONBLOCK))
        return fn(fd, std::forward<Args>(args)...);

    if (-1 == fcntl(fd, F_SETFL, flags | O_NONBLOCK))
        return fn(fd, std::forward<Args>(args)...);

    DebugPrint(dbg_hook, "task(%s) real hook %s fd=%d", tk->DebugInfo(), hook_fn_name, fd);
    ssize_t n = fn(fd, std::forward<Args>(args)...);
    if (n == -1 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
        // get timeout option.
        int timeout_ms = -1;
        struct timeval timeout;
        socklen_t timeout_blen = sizeof(timeout);
        if (0 == getsockopt(fd, SOL_SOCKET, timeout_so, &timeout, &timeout_blen)) {
            if (timeout.tv_sec > 0 || timeout.tv_usec > 0) {
                timeout_ms = timeout.tv_sec * 1000 + timeout.tv_usec / 1000;
                DebugPrint(dbg_hook, "hook task(%s) %s timeout=%dms. fd=%d",
                        g_Scheduler.GetCurrentTaskDebugInfo(), hook_fn_name, timeout_ms, fd);
            }
        }
        auto start_time = std::chrono::system_clock::now();

retry:
        // add into epoll, and switch other context.
        g_Scheduler.IOBlockSwitch(fd, event, timeout_ms);
        bool is_timeout = false;
        if (tk->GetIoWaitData().io_block_timer_) {
            is_timeout = true;
            if (g_Scheduler.BlockCancelTimer(tk->GetIoWaitData().io_block_timer_)) {
                is_timeout = false;
                tk->DecrementRef(); // timer use ref.
            }
        }

        if (tk->GetIoWaitData().wait_successful_ == 0) {
            if (is_timeout) {
                fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
                errno = EAGAIN;
                return -1;
            } else {
                fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
                return fn(fd, std::forward<Args>(args)...);
            }
        }

        DebugPrint(dbg_hook, "continue task(%s) %s. fd=%d", g_Scheduler.GetCurrentTaskDebugInfo(), hook_fn_name, fd);
        n = fn(fd, std::forward<Args>(args)...);
        if (n == -1 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
            /// epoll误唤醒
            if (timeout_ms == -1)   // 不超时, 继续重试
                goto retry;

            int delay_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                    std::chrono::system_clock::now() - start_time).count();
            if (delay_ms < timeout_ms) {    // 还有剩余的超时时间, 重试一次
                timeout_ms -= delay_ms;
                goto retry;
            }
        }
    } else {
        DebugPrint(dbg_hook, "task(%s) syscall(%s) completed immediately. fd=%d",
                g_Scheduler.GetCurrentTaskDebugInfo(), hook_fn_name, fd);
    }

    int e = errno;
    fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
    errno = e;
    return n;
}

int connect(int fd, const struct sockaddr *addr, socklen_t addrlen)
{
    Task* tk = g_Scheduler.GetCurrentTask();
    DebugPrint(dbg_hook, "task(%s) hook connect. %s coroutine.",
            tk ? tk->DebugInfo() : "nil", g_Scheduler.IsCoroutine() ? "In" : "Not in");

    if (!tk) {
        return connect_f(fd, addr, addrlen);
    } else {
        int flags = fcntl(fd, F_GETFL, 0);
        if (flags & O_NONBLOCK)
            return connect_f(fd, addr, addrlen);

        if (-1 == fcntl(fd, F_SETFL, flags | O_NONBLOCK))
            return connect_f(fd, addr, addrlen);

        int n = connect_f(fd, addr, addrlen);
        int e = errno;
        if (n == 0) {
            fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
            DebugPrint(dbg_hook, "continue task(%s) connect completed immediately. fd=%d",
                    g_Scheduler.GetCurrentTaskDebugInfo(), fd);
            return 0;
        } else if (n != -1 || errno != EINPROGRESS) {
            fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
            errno = e;
            return n;
        } else {
            // add into epoll, and switch other context.
            g_Scheduler.IOBlockSwitch(fd, EPOLLOUT, -1);
        }

        if (tk->GetIoWaitData().wait_successful_ == 0) {
            // 添加到epoll中失败了
            fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
            errno = e;
            return n;
        }

        DebugPrint(dbg_hook, "continue task(%s) connect. fd=%d", g_Scheduler.GetCurrentTaskDebugInfo(), fd);
        int error = 0;
        socklen_t len = sizeof(int);
        if (0 == getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &len)) {
            if (0 == error) {
                fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
                DebugPrint(dbg_hook, "continue task(%s) connect success async. fd=%d",
                        g_Scheduler.GetCurrentTaskDebugInfo(), fd);
                return 0;
            } else {
                fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
                errno = error;
                return -1;
            }
        }

        e = errno;      // errno set by getsockopt.
        fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
        errno = e;
        return -1;
    }
}

int IoWait::WaitLoop(bool enable_block)
{
    int c = 0;
    for (;;) {
        std::list<CoTimerPtr> timers;
        timer_mgr_.GetExpired(timers, 128);
        if (timers.empty())
            break;

        c += timers.size();
        // 此处暂存callback而不是Task*，是为了block_cancel能够真实有效。
        std::unique_lock<LFLock> lock(timeout_list_lock_);
        timeout_list_.merge(std::move(timers));
    }

    std::unique_lock<LFLock> lock(epoll_lock_, std::defer_lock);
    if (!lock.try_lock())
        return c ? c : -1;

    ++loop_index_;
    int epoll_n = 0;
    if (IsEpollCreated())
    {
        static epoll_event *evs = new epoll_event[epoll_event_size_];
        for (int epoll_type = 0; epoll_type < 2; ++epoll_type)
        {
retry:
            int timeout = (enable_block && epoll_type == (int)EpollType::read && !c) ? epollwait_ms_ : 0;
            int n = epoll_wait(GetEpoll(epoll_type), evs, epoll_event_size_, timeout);
            if (n == -1) {
                if (errno == EINTR) {
                    goto retry;
                }
                continue;
            }

            epoll_n += n;
            DebugPrint(dbg_scheduler, "do epoll(%d) event, n = %d", epoll_type, n);
            for (int i = 0; i < n; ++i)
            {
                EpollPtr* ep = (EpollPtr*)evs[i].data.ptr;
                ep->revent = evs[i].events;
                Task* tk = ep->tk;
                ++tk->GetIoWaitData().wait_successful_;
                // 将tk暂存, 最后再执行Cancel, 是为了poll和select可以得到正确的计数。
                // 以防Task被加入runnable列表后，被其他线程执行
                epollwait_tasks_.insert(EpollWaitSt{tk, ep->io_block_id});
                DebugPrint(dbg_ioblock,
                        "task(%s) epoll(%s) trigger fd=%d io_block_id(%u) ep(%p) loop_index(%llu)",
                        tk->DebugInfo(), EpollTypeName(epoll_type),
                        ep->fdst->fd, ep->io_block_id, ep, (unsigned long long)loop_index_);
            }
        }

        for (auto &st : epollwait_tasks_)
            Cancel(st.tk, st.id);
        epollwait_tasks_.clear();
    }

    std::list<CoTimerPtr> timeout_list;
    {
        std::unique_lock<LFLock> lock(timeout_list_lock_);
        timeout_list_.swap(timeout_list);
    }

    for (auto &cb : timeout_list)
        (*cb)();

    // 由于epoll_wait的结果中会残留一些未计数的Task*,
    //     epoll的性质决定了这些Task无法计数, 
    //     所以这个析构的操作一定要在epoll_lock的保护中做
    std::vector<SList<Task>> delete_lists;
    Task::PopDeleteList(delete_lists);
    for (auto &delete_list : delete_lists)
        for (auto it = delete_list.begin(); it != delete_list.end();)
        {
            Task* tk = &*it++;
            DebugPrint(dbg_task, "task(%s) delete.", tk->DebugInfo());
            delete tk;
        }

    return epoll_n + c;
}