C++20协程调度框架¶

协程调度框架为异步逻辑提供高性能且易于维护的架构范式。应用场景包括：为每个交易标的创建独立协程以处理高速行情流；为每笔订单创建协程来追踪其从下单到成交的完整生命周期；以及将复杂的交易策略（如套利或做市）封装成状态清晰的协程任务。

#pragma once

#include <atomic>

// For _mm_pause intrinsic
#if defined(_MSC_VER) || defined(__INTEL_COMPILER)
#include <intrin.h>
#define SPIN_PAUSE() _mm_pause()
#elif defined(__GNUC__) || defined(__clang__)
#if defined(__x86_64__) || defined(__i386__)
#include <x86intrin.h>
#define SPIN_PAUSE() _mm_pause()
#else
#define SPIN_PAUSE() ((void)0)
#endif
#else
#define SPIN_PAUSE() ((void)0)
#endif

namespace MyScheduler {

class SpinLock {
 public:
  SpinLock() : flag_(ATOMIC_FLAG_INIT) {}

  void lock() noexcept {
    while (flag_.test_and_set(std::memory_order_acquire)) {
      SPIN_PAUSE();
    }
  }

  void unlock() noexcept { flag_.clear(std::memory_order_release); }

 private:
  std::atomic_flag flag_;
};

class SpinLockGuard {
 public:
  explicit SpinLockGuard(SpinLock& lock) : lock_(lock) { lock_.lock(); }
  ~SpinLockGuard() { lock_.unlock(); }
  SpinLockGuard(const SpinLockGuard&) = delete;
  SpinLockGuard& operator=(const SpinLockGuard&) = delete;

 private:
  SpinLock& lock_;
};

}  // namespace MyScheduler

#pragma once

#include <coroutine>
#include <optional>
#include <type_traits>
#include <utility>

namespace MyScheduler {

using CoroHandle = std::coroutine_handle<>;

template <typename T = void>
class Task {
 public:
  struct promise_type {
    std::optional<T> value;
    auto get_return_object() { return Task{CoroHandle::from_promise(*this)}; }
    auto initial_suspend() { return std::suspend_always{}; }
    auto final_suspend() noexcept { return std::suspend_always{}; }
    void unhandled_exception() { std::terminate(); }

    template <typename U,
              typename = std::enable_if_t<std::is_convertible_v<U, T>>>
    void return_value(U&& val) {
      value = std::forward<U>(val);
    }
  };
  CoroHandle handle;
};

template <>
class Task<void> {
 public:
  struct promise_type {
    auto get_return_object() { return Task{CoroHandle::from_promise(*this)}; }
    auto initial_suspend() { return std::suspend_always{}; }
    auto final_suspend() noexcept { return std::suspend_always{}; }
    void unhandled_exception() { std::terminate(); }
    void return_void() {}
  };
  CoroHandle handle;
};

}  // namespace MyScheduler

#pragma once

#include <atomic>
#include <cstdint>
#include <iostream>
#include <thread>
#include <vector>

#include "concurrentqueue.h"
#include "task.hpp"

#if defined(__linux__)
#include <pthread.h>
#endif

namespace MyScheduler {

using WorkerID = uint32_t;
using CpuID = uint32_t;

class Worker {
 public:
  Worker(WorkerID id, CpuID cpu_id,
         moodycamel::ConcurrentQueue<CoroHandle>& task_queue,
         std::atomic<size_t>& total_tasks, std::atomic<size_t>& finished_tasks,
         uint32_t local_task_capacity = 4);
  ~Worker();

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

  void start();
  void join();

 private:
  void run();
  void set_affinity();

  WorkerID id_;
  CpuID cpu_id_;
  std::thread thread_;

  moodycamel::ConcurrentQueue<CoroHandle>& task_queue_;
  std::atomic<size_t>& total_task_count_;
  std::atomic<size_t>& total_finish_count_;

  const uint32_t local_task_capacity_;
  std::vector<CoroHandle> local_tasks_;
};

// --- Implementation ---

inline Worker::Worker(WorkerID id, CpuID cpu_id,
                      moodycamel::ConcurrentQueue<CoroHandle>& task_queue,
                      std::atomic<size_t>& total_tasks,
                      std::atomic<size_t>& finished_tasks,
                      uint32_t local_task_capacity)
    : id_(id),
      cpu_id_(cpu_id),
      task_queue_(task_queue),
      total_task_count_(total_tasks),
      total_finish_count_(finished_tasks),
      local_task_capacity_(local_task_capacity),
      local_tasks_(local_task_capacity) {}

inline Worker::~Worker() { join(); }

inline void Worker::start() { thread_ = std::thread(&Worker::run, this); }

inline void Worker::join() {
  if (thread_.joinable()) {
    thread_.join();
  }
}

inline void Worker::run() {
  set_affinity();
  uint32_t current_slot = 0;

  while (true) {
    auto& handle = local_tasks_[current_slot];
    if (!handle) {
      if (!task_queue_.try_dequeue(handle)) {
        if (total_finish_count_.load(std::memory_order_acquire) ==
            total_task_count_.load(std::memory_order_acquire)) {
          break;
        }
        std::this_thread::yield();
      }
    }
    if (handle) {
      handle.resume();
      if (handle.done()) {
        handle = nullptr;
        total_finish_count_.fetch_add(1, std::memory_order_release);
      }
    }
    current_slot = (current_slot + 1) % local_task_capacity_;
  }
}

inline void Worker::set_affinity() {
#if defined(__linux__)
  cpu_set_t cpuset;
  CPU_ZERO(&cpuset);
  CPU_SET(cpu_id_, &cpuset);
  if (pthread_setaffinity_np(thread_.native_handle(), sizeof(cpu_set_t),
                             &cpuset) != 0) {
    std::cerr << "Error setting thread affinity for worker " << id_
              << " on CPU " << cpu_id_ << std::endl;
  }
#endif
}

}  // namespace MyScheduler

#pragma once

#include <atomic>
#include <cstdint>
#include <memory>
#include <vector>

#include "concurrentqueue.h"
#include "spinlock.hpp"
#include "task.hpp"
#include "worker.hpp"

namespace MyScheduler {

class Scheduler {
 public:
  explicit Scheduler(const std::vector<CpuID>& cpu_ids);
  ~Scheduler();

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

  void start();
  void stop();
  void schedule(CoroHandle handle);

  class Awaiter {
   public:
    explicit Awaiter(Scheduler& scheduler) noexcept : scheduler_(scheduler) {}
    bool await_ready() const noexcept { return false; }
    void await_suspend(CoroHandle awaiting_coroutine) noexcept {
      scheduler_.task_queue_->enqueue(awaiting_coroutine);
    }
    void await_resume() noexcept {}

   private:
    Scheduler& scheduler_;
  };

  auto yield() noexcept { return Awaiter{*this}; }

 private:
  std::vector<CpuID> cpus_;
  std::vector<std::unique_ptr<Worker>> workers_;

  std::atomic<size_t> total_task_count_{0};
  std::atomic<size_t> total_finish_count_{0};

  SpinLock enqueue_lock_;
  std::unique_ptr<moodycamel::ConcurrentQueue<CoroHandle>> task_queue_;
};

// --- Implementation ---

inline Scheduler::Scheduler(const std::vector<CpuID>& cpu_ids)
    : cpus_(cpu_ids),
      task_queue_(std::make_unique<moodycamel::ConcurrentQueue<CoroHandle>>()) {
}

inline Scheduler::~Scheduler() { stop(); }

inline void Scheduler::start() {
  workers_.reserve(cpus_.size());
  for (size_t i = 0; i < cpus_.size(); ++i) {
    workers_.emplace_back(std::make_unique<Worker>(
        i, cpus_[i], *task_queue_, total_task_count_, total_finish_count_));
  }
  for (auto& worker : workers_) {
    worker->start();
  }
}

inline void Scheduler::stop() {
  for (auto& worker : workers_) {
    worker->join();
  }
  workers_.clear();
}

inline void Scheduler::schedule(CoroHandle handle) {
  if (handle) {
    SpinLockGuard lock(enqueue_lock_);
    total_task_count_.fetch_add(1, std::memory_order_relaxed);
    task_queue_->enqueue(handle);
  }
}

}  // namespace MyScheduler

#include <iostream>
#include <thread>
#include <vector>

#include "scheduler.hpp"

MyScheduler::Scheduler* g_scheduler = nullptr;

MyScheduler::Task<> coroutine_task(int id, int iterations) {
  for (int i = 0; i < iterations; ++i) {
    std::cout << "Task " << id << " on iteration " << i
              << " is running on thread " << std::this_thread::get_id()
              << std::endl;
    co_await g_scheduler->yield();
  }
  std::cout << "Task " << id << " finished." << std::endl;
}

int main() {
  unsigned int num_cpus = std::thread::hardware_concurrency();
  if (num_cpus == 0) num_cpus = 4;
  std::cout << "Starting scheduler with " << num_cpus << " worker threads."
            << std::endl;

  std::vector<MyScheduler::CpuID> cpu_ids(num_cpus);
  for (unsigned int i = 0; i < num_cpus; ++i) cpu_ids[i] = i;

  MyScheduler::Scheduler scheduler(cpu_ids);
  g_scheduler = &scheduler;
  scheduler.start();

  const int num_tasks = 10;
  std::vector<MyScheduler::Task<>> tasks;
  tasks.reserve(num_tasks);
  for (int i = 0; i < num_tasks; ++i) {
    tasks.push_back(coroutine_task(i, 5));
  }

  for (auto& task : tasks) {
    scheduler.schedule(task.handle);
  }

  scheduler.stop();
  std::cout << "All tasks completed. Scheduler shut down." << std::endl;
  return 0;
}