推理请求动态批处理调度器（动态 batch）

/**
题目：推理请求动态批处理调度器（简化版）

**题目描述**：实现一个批处理调度器，用于推理服务收包后聚合请求成 batch 送 GPU：
- 请求以 `submit(req)` 不断到达，`req` 含 `id` 和 `arrive_time_ms`。
- 调度器输出 batch 列表（一个 batch 是请求 id 的数组），触发规则（到达即触发 + 动态超时）：
  1. 一旦队列中累积 `max_batch_size` 个请求，**立即**切出一个 batch 输出；
  2. 否则，当**最早未成批请求**的等待时间达到 `max_wait_ms` 时，把当前已到达的所有请求（≥1 个）切成一个 batch 输出；
- 提供接口：`submit(req)` 和 `poll_batches(now_ms)`（或在 submit 内部判断，二选一均可），`now_ms` 由外部时钟驱动。
**输入输出**：给一组 `(arrive_time_ms, id)` 序列和参数 `max_batch_size=4, max_wait_ms=50`，输出每次 batch 的
`(fire_time_ms, [ids])`。样例：到达时刻 0,10,20,100,110,200 → batch1@50:[0,1,2]（id0 等到 50ms 超时，仅 3 个未满批）、
batch2@150:[3,4]（id3 等到 150ms 超时）、batch3@250:[5]（id5 单独成批）。
*/

#include <cstdint>
#include <deque>
#include <iostream>
#include <vector>

// 单个请求。
struct Request {
    int id;
    std::int64_t arrive_time_ms;
};

struct Batch {
    std::int64_t fire_time_ms;  // 触发时间
    std::vector<int> ids;
};

// 只负责调度，实际的定时器触发由外部时钟驱动。
// Note：外部来保证请求的 arrive_time_ms 是单调递增的，外部时钟驱动 poll_batches(now_ms) 也是单调递增的。
class DynamicBatchScheduler {
public:
    DynamicBatchScheduler(std::size_t max_batch_size, std::int64_t max_wait_ms)
        : max_batch_size_(max_batch_size),
          max_wait_ms_(max_wait_ms),
          current_time_ms_(0) {
        // 满批至少 1 个请求，避免切出空 batch。
        if (max_batch_size_ <= 0) {
            max_batch_size_ = 1;
        }
    }

    // 提交一个请求到调度器。外部保证时间单调递增。只负责提交和触发 batch。
    // 按照 req 时间，超时和满 batch 都会触发规则组 batch
    void submit(const Request& req) {
        if (req.arrive_time_ms < current_time_ms_) {
            std::cerr << "Request arrive_time_ms is earlier than current_time_ms_" << std::endl;
            return;
        }

        // 外部时钟推进到该请求的到达时刻，先冲刷已超时的请求（规则 2）。
        current_time_ms_ = req.arrive_time_ms;
        updateTimeoutReq();

        // 入队后若满批，立即切出一个 batch（规则 1）。
        queue_.push_back(req);
        if (queue_.size() >= max_batch_size_) {
            batches_.push_back(createBatch(req.arrive_time_ms, max_batch_size_));
        }
    }

    // 外部时钟推进到 now_ms，冲刷所有已超时的请求（规则 2），
    // 并把当前已完成、待取走的 batch 取出并清空后返回（drain 语义）。
    std::vector<Batch> poll_batches(std::int64_t now_ms) {
        current_time_ms_ = now_ms;
        updateTimeoutReq();

        std::vector<Batch> out;
        out.swap(batches_);
        return out;
    }

private:
    // 规则 2：从队首开始，把等待时间已达 max_wait_ms 的请求连续切出。
    void updateTimeoutReq() {
        while (!queue_.empty()) {
            const std::int64_t earliest = queue_.front().arrive_time_ms;
            if (current_time_ms_ < earliest + max_wait_ms_) {
                break;
            }
            // 超时触发：当前已到达的所有请求（≥1）一起切出。
            batches_.push_back(createBatch(earliest + max_wait_ms_, queue_.size()));
        }
    }

    // 从队首切出 count 个请求作为一个 batch。
    Batch createBatch(std::int64_t fire_time_ms, std::size_t count) {
        Batch b;
        b.fire_time_ms = fire_time_ms;
        for (std::size_t i = 0; i < count && !queue_.empty(); ++i) {
            b.ids.push_back(queue_.front().id);
            queue_.pop_front();
        }
        return b;
    }

    std::size_t max_batch_size_;
    std::int64_t max_wait_ms_;
    std::int64_t current_time_ms_;

    std::deque<Request> queue_;
    std::vector<Batch> batches_;
};

// 跑一个用例：提交所有请求，再把时钟推进到 final_now_ms 冲刷残余，打印结果。
static void run_case(const char* name,
                     std::size_t max_batch_size,
                     std::int64_t max_wait_ms,
                     const std::vector<Request>& requests,
                     std::int64_t final_now_ms) {
    DynamicBatchScheduler sched(max_batch_size, max_wait_ms);
    for (const Request& r : requests) {
        sched.submit(r);
    }
    const std::vector<Batch> batches = sched.poll_batches(final_now_ms);

    std::cout << name << ":\n";
    for (const Batch& b : batches) {
        std::cout << "  batch@" << b.fire_time_ms << ": [";
        for (std::size_t i = 0; i < b.ids.size(); ++i) {
            if (i != 0) std::cout << ",";
            std::cout << b.ids[i];
        }
        std::cout << "]\n";
    }
}

int main() {
    // 样例：到达时刻 0,10,20,100,110,200。
    run_case("case_timeout", 4, 50, {{0, 0}, {1, 10}, {2, 20}, {3, 100}, {4, 110}, {5, 200}}, 200);

    // 满批用例：到达时刻 0,10,20,30，凑满 4 个立即触发（规则 1）。
    run_case("case_full", 4, 50, {{0, 0}, {1, 10}, {2, 20}, {3, 30}}, 30);

    return 0;
}

class RealtimeBatchScheduler {
public:
    using Callback = std::function<void(const Batch&)>;

    RealtimeBatchScheduler(std::size_t max_batch_size, std::int64_t max_wait_ms,
                           std::chrono::milliseconds poll_interval = std::chrono::milliseconds(5))
        : core_(max_batch_size, max_wait_ms),
          poll_interval_(poll_interval),
          stop_(false) {
        timer_ = std::thread(&RealtimeBatchScheduler::run, this);
    }

    ~RealtimeBatchScheduler() { stop(); }

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

    void set_callback(Callback cb) {
        std::lock_guard<std::mutex> lk(mtx_);
        callback_ = std::move(cb);
    }

    // 请求到达：盖真实时间戳 → 交给 core → 立即 drain（满批/已到期超时）并回调。
    void submit(const Request& req) {
        std::vector<Batch> fired;
        {
            std::lock_guard<std::mutex> lk(mtx_);
            Request r = req;
            r.arrive_time_ms = now_ms();
            core_.submit(r);
            fired = core_.poll_batches(r.arrive_time_ms);
        }
        for (const Batch& b : fired) {
            deliver(b);
        }
    }

    void stop() {
        {
            std::lock_guard<std::mutex> lk(mtx_);
            if (stop_) {
                return;
            }
            stop_ = true;
        }
        cv_.notify_all();
        if (timer_.joinable()) {
            timer_.join();
        }
    }

private:
    using Clock = std::chrono::steady_clock;

    static std::int64_t now_ms() {
        return std::chrono::duration_cast<std::chrono::milliseconds>(
                   Clock::now().time_since_epoch())
            .count();
    }

    // 周期轮询：每隔 poll_interval_ 推进一次时钟，drain 超时 batch。
    void run() {
        std::unique_lock<std::mutex> lk(mtx_);
        while (true) {
            if (stop_) {
                return;
            }
            if (cv_.wait_for(lk, poll_interval_, [this] { return stop_; })) {
                return;  // 被 stop 唤醒
            }
            std::vector<Batch> fired = core_.poll_batches(now_ms());
            lk.unlock();
            for (const Batch& b : fired) {
                deliver(b);
            }
            lk.lock();
        }
    }

    // 先拷贝回调再在锁外调用，避免回调里再进 submit 死锁。
    void deliver(const Batch& b) {
        Callback cb;
        {
            std::lock_guard<std::mutex> lk(mtx_);
            cb = callback_;
        }
        if (cb) {
            cb(b);
        }
    }

    DynamicBatchScheduler core_;
    std::chrono::milliseconds poll_interval_;
    bool stop_;

    std::mutex mtx_;
    std::condition_variable cv_;
    Callback callback_;
    std::thread timer_;
};

// 真实定时器演示：稀疏到达触发超时，密集到达触发满批。
static void demo_realtime() {
    const auto start = std::chrono::steady_clock::now();
    const std::int64_t start_ms =
        std::chrono::duration_cast<std::chrono::milliseconds>(start.time_since_epoch()).count();

    RealtimeBatchScheduler sched(/*max_batch_size=*/4, /*max_wait_ms=*/50);
    sched.set_callback([start_ms](const Batch& b) {
        std::cout << "  batch@" << (b.fire_time_ms - start_ms) << ": [";
        for (std::size_t i = 0; i < b.ids.size(); ++i) {
            if (i != 0) std::cout << ",";
            std::cout << b.ids[i];
        }
        std::cout << "]\n";
    });

    std::cout << "realtime:\n";

    // 稀疏：3 个（不满批），id0 会等到 50ms 超时被切出。
    sched.submit({0, 0});
    std::this_thread::sleep_for(std::chrono::milliseconds(10));
    sched.submit({1, 0});
    std::this_thread::sleep_for(std::chrono::milliseconds(10));
    sched.submit({2, 0});
    std::this_thread::sleep_for(std::chrono::milliseconds(100));  // 等超时触发

    // 密集：连发 4 个，凑满立即触发。
    sched.submit({3, 0});
    sched.submit({4, 0});
    sched.submit({5, 0});
    sched.submit({6, 0});
    std::this_thread::sleep_for(std::chrono::milliseconds(100));  // 等剩余超时触发

    sched.stop();
}

/**
题目：推理请求动态批处理调度器 —— 三层设计版（真实定时器）

分层（从 0 开始的推荐结构）：
1. Request / Batch —— 纯数据。
2. BatchEngine —— 纯策略引擎（组 batch 规则），无线程、无真实时钟、无锁。
   只回答"给定当前时间，哪些请求该成批"。submit / advance_to 直接返回触发的 batch。
3. RealtimeBatchScheduler —— 真实时钟适配器，唯一碰线程/锁/真实时钟的地方。
   持有一个 BatchEngine，起一个线程用 condition_variable 等到队首 deadline 自动触发超时。

用法：
- 确定性测试：直接驱动 BatchEngine（submit/advance_to 传显式时间），零抖动。
- 生产：用 RealtimeBatchScheduler，submit 到达 + set_callback 收 batch。
*/

#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <deque>
#include <functional>
#include <iostream>
#include <mutex>
#include <thread>
#include <vector>

// ===== 第 1 层：纯数据 =====
struct Request {
    int id;
    std::int64_t arrive_time_ms;
};

struct Batch {
    std::int64_t fire_time_ms;  // 触发时间
    std::vector<int> ids;
};

// ===== 第 2 层：纯策略引擎 =====
// 单线程、无真实时钟、无 I/O。组 batch 的规则只在这里有一份。
class BatchEngine {
public:
    BatchEngine(std::size_t max_batch_size, std::int64_t max_wait_ms)
        : max_batch_size_(max_batch_size ? max_batch_size : 1),
          max_wait_ms_(max_wait_ms),
          current_time_ms_(0) {}

    // 到达一个请求（按到达时间顺序调用）。返回因此触发的 batch（超时 + 满批）。
    std::vector<Batch> submit(const Request& req) {
        if (req.arrive_time_ms < current_time_ms_) {
            std::cerr << "arrive_time_ms is earlier than current_time_ms_" << std::endl;
            return {};
        }

        // 外部时钟推进到该请求的到达时刻，先冲刷已超时的请求（规则 2）。
        current_time_ms_ = req.arrive_time_ms;
        std::vector<Batch> fired = flushTimeouts();

        // 入队后若满批，立即切出一个 batch（规则 1）。
        queue_.push_back(req);
        if (queue_.size() >= max_batch_size_) {
            fired.push_back(cut(req.arrive_time_ms, max_batch_size_));
        }
        return fired;
    }

    // 推进时钟到 now_ms，返回因此触发的超时 batch。
    std::vector<Batch> advance_to(std::int64_t now_ms) {
        current_time_ms_ = now_ms;
        return flushTimeouts();
    }

    // 队首请求的 deadline（到达 + max_wait）；空队列返回 -1。
    std::int64_t next_deadline_ms() const {
        return queue_.empty() ? -1 : queue_.front().arrive_time_ms + max_wait_ms_;
    }

private:
    // 规则 2：最早请求等待达 max_wait_ms，把当前所有请求一起切出。
    std::vector<Batch> flushTimeouts() {
        std::vector<Batch> fired;
        while (!queue_.empty()) {
            const std::int64_t earliest = queue_.front().arrive_time_ms;
            if (current_time_ms_ < earliest + max_wait_ms_) {
                break;
            }
            fired.push_back(cut(earliest + max_wait_ms_, queue_.size()));
        }
        return fired;
    }

    // 从队首切出 count 个请求为一个 batch。
    Batch cut(std::int64_t fire_time_ms, std::size_t count) {
        Batch b;
        b.fire_time_ms = fire_time_ms;
        for (std::size_t i = 0; i < count && !queue_.empty(); ++i) {
            b.ids.push_back(queue_.front().id);
            queue_.pop_front();
        }
        return b;
    }

    std::size_t max_batch_size_;
    std::int64_t max_wait_ms_;
    std::int64_t current_time_ms_;
    std::deque<Request> queue_;
};

// ===== 第 3 层：真实时钟适配器 =====
// 引擎由一把 mutex 串行化所有访问（引擎本身不加锁）。定时线程用 condition_variable
// 等到队首 deadline，到点推进引擎并回调超时 batch；满批在 submit 时立即回调。
class RealtimeBatchScheduler {
public:
    using Callback = std::function<void(const Batch&)>;

    RealtimeBatchScheduler(std::size_t max_batch_size, std::int64_t max_wait_ms)
        : engine_(max_batch_size, max_wait_ms), stop_(false) {
        timer_ = std::thread(&RealtimeBatchScheduler::run, this);
    }

    ~RealtimeBatchScheduler() { stop(); }

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

    void set_callback(Callback cb) {
        std::lock_guard<std::mutex> lk(mtx_);
        callback_ = std::move(cb);
    }

    // 请求到达：盖真实时间戳，交给引擎，触发的 batch 通过回调交付。
    void submit(const Request& req) {
        std::vector<Batch> fired;
        {
            std::lock_guard<std::mutex> lk(mtx_);
            Request r = req;
            r.arrive_time_ms = now_ms();
            fired = engine_.submit(r);
        }
        for (const Batch& b : fired) {
            deliver(b);
        }
        cv_.notify_all();  // 唤醒定时线程重算 deadline
    }

    void stop() {
        {
            std::lock_guard<std::mutex> lk(mtx_);
            if (stop_) {
                return;
            }
            stop_ = true;
        }
        cv_.notify_all();
        if (timer_.joinable()) {
            timer_.join();
        }
    }

private:
    using Clock = std::chrono::steady_clock;
    using TimePoint = Clock::time_point;

    static std::int64_t now_ms() {
        return std::chrono::duration_cast<std::chrono::milliseconds>(
                   Clock::now().time_since_epoch())
            .count();
    }

    static TimePoint from_ms(std::int64_t ms) {
        return TimePoint(std::chrono::milliseconds(ms));
    }

    // 定时线程：等到队首 deadline，到点推进引擎并回调超时 batch。
    void run() {
        std::unique_lock<std::mutex> lk(mtx_);
        while (true) {
            if (stop_) {
                return;
            }

            const std::int64_t deadline = engine_.next_deadline_ms();
            if (deadline < 0) {  // 空队列
                cv_.wait(lk, [this] { return stop_ || engine_.next_deadline_ms() >= 0; });
                continue;
            }

            // 等到 deadline；若期间队首变化 / 停止，则提前醒来重算。
            const bool timed_out = !cv_.wait_until(lk, from_ms(deadline), [this, deadline] {
                return stop_ || engine_.next_deadline_ms() != deadline;
            });

            if (stop_) {
                return;
            }
            if (timed_out) {
                std::vector<Batch> fired = engine_.advance_to(now_ms());
                lk.unlock();
                for (const Batch& b : fired) {
                    deliver(b);
                }
                lk.lock();
            }
        }
    }

    // 先拷贝回调再在锁外调用，避免回调里再进 submit 死锁。
    void deliver(const Batch& b) {
        Callback cb;
        {
            std::lock_guard<std::mutex> lk(mtx_);
            cb = callback_;
        }
        if (cb) {
            cb(b);
        }
    }

    BatchEngine engine_;
    bool stop_;

    std::mutex mtx_;
    std::condition_variable cv_;
    Callback callback_;
    std::thread timer_;
};

// 打印一批 batch。
static void print_batches(const char* name, const std::vector<Batch>& batches) {
    std::cout << name << ":\n";
    for (const Batch& b : batches) {
        std::cout << "  batch@" << b.fire_time_ms << ": [";
        for (std::size_t i = 0; i < b.ids.size(); ++i) {
            if (i != 0) std::cout << ",";
            std::cout << b.ids[i];
        }
        std::cout << "]\n";
    }
}

int main() {
    // ===== 确定性：直接驱动引擎（零线程、零抖动，可做单测） =====
    {
        BatchEngine engine(/*max_batch_size=*/4, /*max_wait_ms=*/50);
        const std::vector<Request> reqs = {
            {0, 0}, {1, 10}, {2, 20}, {3, 100}, {4, 110}, {5, 200},
        };
        std::vector<Batch> all;
        for (const Request& r : reqs) {
            std::vector<Batch> f = engine.submit(r);
            all.insert(all.end(), f.begin(), f.end());
        }
        std::vector<Batch> tail = engine.advance_to(/*now_ms=*/250);
        all.insert(all.end(), tail.begin(), tail.end());
        print_batches("engine_timeout", all);
    }
    {
        BatchEngine engine(/*max_batch_size=*/4, /*max_wait_ms=*/50);
        std::vector<Batch> all;
        for (const Request& r : std::vector<Request>{{0, 0}, {1, 10}, {2, 20}, {3, 30}}) {
            std::vector<Batch> f = engine.submit(r);
            all.insert(all.end(), f.begin(), f.end());
        }
        print_batches("engine_full", all);
    }

    // ===== 真实定时器：自动触发超时 + 满批 =====
    {
        RealtimeBatchScheduler sched(/*max_batch_size=*/4, /*max_wait_ms=*/50);

        const auto start = std::chrono::steady_clock::now();
        auto rel_ms = [&start] {
            return std::chrono::duration_cast<std::chrono::milliseconds>(
                       std::chrono::steady_clock::now() - start)
                .count();
        };

        sched.set_callback([&rel_ms](const Batch& b) {
            std::cout << "  [" << rel_ms() << "ms] batch: [";
            for (std::size_t i = 0; i < b.ids.size(); ++i) {
                if (i != 0) std::cout << ",";
                std::cout << b.ids[i];
            }
            std::cout << "]\n";
        });

        std::cout << "realtime:\n";

        // 稀疏：3 个（不满批），id0 会等到 50ms 超时被切出。
        // 注意：真实版会盖真实时间戳，第二个字段（arrive_time_ms）用 0 占位即可。
        sched.submit({0, 0});
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
        sched.submit({1, 0});
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
        sched.submit({2, 0});
        std::this_thread::sleep_for(std::chrono::milliseconds(100));  // 等超时触发

        // 密集：连发 4 个，凑满立即触发。
        sched.submit({3, 0});
        sched.submit({4, 0});
        sched.submit({5, 0});
        sched.submit({6, 0});
        std::this_thread::sleep_for(std::chrono::milliseconds(100));  // 等剩余超时触发

        sched.stop();
    }
    return 0;
}