swarm_engine_core/learn/daemon/

subscriber.rs

//! Event Subscribers - Event を受信して Record に変換し LearningDaemon に送信
//!
//! ## 設計原則: Event ごとに Subscriber を追加
//!
//! **新しい Event Channel を追加したら、対応する Subscriber も追加すること。**
//!
//! 各 Subscriber は単一責務（Single Responsibility）で設計されている：
//! - `ActionEventSubscriber`: `ActionEvent` を処理
//! - `LearningEventSubscriber`: `LearningEvent` を処理
//!
//! これにより：
//! - 各 Subscriber のテストが独立
//! - 新しい Event 追加時の影響範囲が限定的
//! - 必要な Subscriber のみを起動可能
//!
//! ## Event → Record → Episode の流れ
//!
//! **Subscriber は Event と Record の橋渡し役。**
//! Production でオンラインで学習データを収集するための重要なコンポーネント。
//!
//! ```text
//! [Event Layer]              [Subscriber]                    [Learn Layer]
//!
//! ActionEventPublisher ──▶ ActionEventSubscriber ──────┐
//!                                                      │
//!                          Record::from(&event)        ├──▶ LearningDaemon
//!                                                      │         │
//! LearningEventChannel ──▶ LearningEventSubscriber ────┘         ▼
//!                                                           Episode
//!                                                              │
//!                                                              ▼
//!                                                         LoRA / Learn
//! ```
//!
//! ## 新しい Event Channel を追加する際（Checklist）
//!
//! 1. `events/` に Event 型を定義
//! 2. `learn/record/` に `From<&Event> for *Record` を実装
//! 3. `record/mod.rs` で `From<&Event> for Record` にルーティング追加
//! 4. **このファイルに対応する `*EventSubscriber` を追加**
//! 5. `mod.rs` で re-export
//! 6. 使用箇所で Subscriber を起動
//!
//! **重要**: Builder で直接 Record を作らない。必ず Event 経由で変換する。
//!
//! ## 対応 Event と Subscriber
//!
//! | Event | Subscriber | Record | 変換場所 |
//! |-------|------------|--------|----------|
//! | `ActionEvent` | `ActionEventSubscriber` | `ActionRecord` | `record/action.rs` |
//! | `LearningEvent::StrategyAdvice` | `LearningEventSubscriber` | `ActionRecord` | `record/mod.rs` |
//! | `LearningEvent::DependencyGraphInference` | `LearningEventSubscriber` | `DependencyGraphRecord` | `record/dependency_graph.rs` |
//!
//! ## 使用例
//!
//! ```ignore
//! use swarm_engine_core::events::{ActionEventPublisher, LearningEventChannel};
//! use swarm_engine_core::learn::daemon::{
//!     ActionEventSubscriber, LearningEventSubscriber, EventSubscriberConfig, LearningDaemon
//! };
//!
//! let (publisher, _rx) = ActionEventPublisher::new(1024);
//! let daemon = LearningDaemon::new(config, trigger)?;
//! let record_tx = daemon.record_sender();
//!
//! // ActionEvent Subscriber 起動
//! let action_sub = ActionEventSubscriber::new(publisher.subscribe(), record_tx.clone());
//! tokio::spawn(action_sub.run());
//!
//! // LearningEvent Subscriber 起動
//! let learning_channel = LearningEventChannel::global();
//! learning_channel.enable();
//! let learning_sub = LearningEventSubscriber::new(learning_channel.subscribe(), record_tx);
//! tokio::spawn(learning_sub.run());
//! ```

use tokio::sync::{broadcast, mpsc};

use crate::events::{ActionEvent, LearningEvent};
use crate::learn::record::Record;

// ============================================================================
// EventSubscriberConfig
// ============================================================================

/// Event Subscriber の共通設定
///
/// 全ての Subscriber で共有される設定。
/// バッチ処理とフラッシュ間隔を制御する。
#[derive(Debug, Clone)]
pub struct EventSubscriberConfig {
    /// バッチサイズ（この数に達したら送信）
    pub batch_size: usize,
    /// フラッシュ間隔（None の場合は batch_size のみで判定）
    pub flush_interval_ms: Option<u64>,
}

impl Default for EventSubscriberConfig {
    fn default() -> Self {
        Self {
            batch_size: 100,
            flush_interval_ms: Some(1000), // 1秒
        }
    }
}

impl EventSubscriberConfig {
    /// 新しい設定を作成
    pub fn new() -> Self {
        Self::default()
    }

    /// バッチサイズを設定
    pub fn batch_size(mut self, size: usize) -> Self {
        self.batch_size = size;
        self
    }

    /// フラッシュ間隔を設定（ミリ秒）
    pub fn flush_interval_ms(mut self, ms: u64) -> Self {
        self.flush_interval_ms = Some(ms);
        self
    }

    /// フラッシュ間隔を無効化
    pub fn no_flush_interval(mut self) -> Self {
        self.flush_interval_ms = None;
        self
    }
}

// ============================================================================
// ActionEventSubscriber
// ============================================================================

/// ActionEvent を受信して LearningDaemon に Record を送信
///
/// `ActionEventPublisher` から `ActionEvent` を受信し、
/// `Record` に変換してバッチで `LearningDaemon` に送信する。
///
/// ## 対応 Event
///
/// - `ActionEvent` → `ActionRecord`
pub struct ActionEventSubscriber {
    /// ActionEvent 受信チャンネル
    rx: broadcast::Receiver<ActionEvent>,
    /// Record 送信チャンネル（LearningDaemon.record_sender()）
    record_tx: mpsc::Sender<Vec<Record>>,
    /// 設定
    config: EventSubscriberConfig,
    /// バッファ
    buffer: Vec<Record>,
}

impl ActionEventSubscriber {
    /// 新しい ActionEventSubscriber を作成
    pub fn new(rx: broadcast::Receiver<ActionEvent>, record_tx: mpsc::Sender<Vec<Record>>) -> Self {
        Self::with_config(rx, record_tx, EventSubscriberConfig::default())
    }

    /// 設定を指定して ActionEventSubscriber を作成
    pub fn with_config(
        rx: broadcast::Receiver<ActionEvent>,
        record_tx: mpsc::Sender<Vec<Record>>,
        config: EventSubscriberConfig,
    ) -> Self {
        let batch_size = config.batch_size;
        Self {
            rx,
            record_tx,
            config,
            buffer: Vec::with_capacity(batch_size),
        }
    }

    /// 受信ループを開始
    ///
    /// `ActionEvent` を受信し、`Record` に変換してバッチ送信する。
    /// チャンネルが閉じられるか、送信先が閉じられたら終了。
    pub async fn run(mut self) {
        tracing::info!(
            batch_size = self.config.batch_size,
            flush_interval_ms = ?self.config.flush_interval_ms,
            "ActionEventSubscriber started"
        );

        if let Some(interval_ms) = self.config.flush_interval_ms {
            self.run_with_flush_interval(interval_ms).await;
        } else {
            self.run_batch_only().await;
        }

        // Flush remaining records
        self.flush().await;

        tracing::info!("ActionEventSubscriber stopped");
    }

    /// フラッシュ間隔ありの受信ループ
    async fn run_with_flush_interval(&mut self, interval_ms: u64) {
        use std::time::Duration;
        use tokio::time::{interval, Instant};

        let mut flush_interval = interval(Duration::from_millis(interval_ms));
        let mut last_flush = Instant::now();

        loop {
            tokio::select! {
                // ActionEvent 受信
                result = self.rx.recv() => {
                    match result {
                        Ok(event) => {
                            self.buffer.push(Record::from(&event));

                            if self.buffer.len() >= self.config.batch_size {
                                if !self.flush().await {
                                    return;
                                }
                                last_flush = Instant::now();
                            }
                        }
                        Err(broadcast::error::RecvError::Closed) => {
                            tracing::debug!("ActionEvent channel closed");
                            return;
                        }
                        Err(broadcast::error::RecvError::Lagged(n)) => {
                            tracing::warn!(lagged = n, "ActionEventSubscriber lagged behind");
                        }
                    }
                }

                // 定期フラッシュ
                _ = flush_interval.tick() => {
                    if !self.buffer.is_empty() && last_flush.elapsed().as_millis() as u64 >= interval_ms {
                        if !self.flush().await {
                            return;
                        }
                        last_flush = Instant::now();
                    }
                }
            }
        }
    }

    /// バッチサイズのみの受信ループ（フラッシュ間隔なし）
    async fn run_batch_only(&mut self) {
        loop {
            match self.rx.recv().await {
                Ok(event) => {
                    self.buffer.push(Record::from(&event));

                    if self.buffer.len() >= self.config.batch_size && !self.flush().await {
                        return;
                    }
                }
                Err(broadcast::error::RecvError::Closed) => {
                    tracing::debug!("ActionEvent channel closed");
                    return;
                }
                Err(broadcast::error::RecvError::Lagged(n)) => {
                    tracing::warn!(lagged = n, "ActionEventSubscriber lagged behind");
                }
            }
        }
    }

    /// バッファを送信
    ///
    /// Returns: 送信成功なら true、チャンネル閉鎖なら false
    async fn flush(&mut self) -> bool {
        if self.buffer.is_empty() {
            return true;
        }

        let records = std::mem::take(&mut self.buffer);
        let count = records.len();

        match self.record_tx.send(records).await {
            Ok(()) => {
                tracing::debug!(count, "Flushed ActionEvent records to LearningDaemon");
                true
            }
            Err(_) => {
                tracing::warn!("LearningDaemon channel closed");
                false
            }
        }
    }
}

// ============================================================================
// LearningEventSubscriber
// ============================================================================

/// LearningEvent を受信して LearningDaemon に Record を送信
///
/// `LearningEventChannel` から `LearningEvent` を受信し、
/// `Record` に変換してバッチで `LearningDaemon` に送信する。
///
/// ## 対応 Event
///
/// - `LearningEvent::StrategyAdvice` → `StrategyAdviceRecord`
/// - `LearningEvent::DependencyGraphInference` → `DependencyGraphRecord`
/// - `LearningEvent::LearnStatsSnapshot` → `LearnStatsRecord`
///
/// ## 使用例
///
/// ```ignore
/// use swarm_engine_core::events::LearningEventChannel;
/// use swarm_engine_core::learn::daemon::{LearningEventSubscriber, EventSubscriberConfig};
///
/// let channel = LearningEventChannel::global();
/// channel.enable();
///
/// let subscriber = LearningEventSubscriber::new(channel.subscribe(), record_tx);
/// tokio::spawn(subscriber.run());
/// ```
pub struct LearningEventSubscriber {
    /// LearningEvent 受信チャンネル
    rx: broadcast::Receiver<LearningEvent>,
    /// Record 送信チャンネル（LearningDaemon.record_sender()）
    record_tx: mpsc::Sender<Vec<Record>>,
    /// 設定
    config: EventSubscriberConfig,
    /// バッファ
    buffer: Vec<Record>,
}

impl LearningEventSubscriber {
    /// 新しい LearningEventSubscriber を作成
    pub fn new(
        rx: broadcast::Receiver<LearningEvent>,
        record_tx: mpsc::Sender<Vec<Record>>,
    ) -> Self {
        Self::with_config(rx, record_tx, EventSubscriberConfig::default())
    }

    /// 設定を指定して LearningEventSubscriber を作成
    pub fn with_config(
        rx: broadcast::Receiver<LearningEvent>,
        record_tx: mpsc::Sender<Vec<Record>>,
        config: EventSubscriberConfig,
    ) -> Self {
        let batch_size = config.batch_size;
        Self {
            rx,
            record_tx,
            config,
            buffer: Vec::with_capacity(batch_size),
        }
    }

    /// 受信ループを開始
    ///
    /// `LearningEvent` を受信し、`Record` に変換してバッチ送信する。
    /// チャンネルが閉じられるか、送信先が閉じられたら終了。
    pub async fn run(mut self) {
        tracing::info!(
            batch_size = self.config.batch_size,
            flush_interval_ms = ?self.config.flush_interval_ms,
            "LearningEventSubscriber started"
        );

        if let Some(interval_ms) = self.config.flush_interval_ms {
            self.run_with_flush_interval(interval_ms).await;
        } else {
            self.run_batch_only().await;
        }

        // Flush remaining records
        self.flush().await;

        tracing::info!("LearningEventSubscriber stopped");
    }

    /// フラッシュ間隔ありの受信ループ
    async fn run_with_flush_interval(&mut self, interval_ms: u64) {
        use std::time::Duration;
        use tokio::time::{interval, Instant};

        let mut flush_interval = interval(Duration::from_millis(interval_ms));
        let mut last_flush = Instant::now();

        loop {
            tokio::select! {
                // LearningEvent 受信
                result = self.rx.recv() => {
                    match result {
                        Ok(event) => {
                            self.buffer.push(Record::from(&event));

                            if self.buffer.len() >= self.config.batch_size {
                                if !self.flush().await {
                                    return;
                                }
                                last_flush = Instant::now();
                            }
                        }
                        Err(broadcast::error::RecvError::Closed) => {
                            tracing::debug!("LearningEvent channel closed");
                            return;
                        }
                        Err(broadcast::error::RecvError::Lagged(n)) => {
                            tracing::warn!(lagged = n, "LearningEventSubscriber lagged behind");
                        }
                    }
                }

                // 定期フラッシュ
                _ = flush_interval.tick() => {
                    if !self.buffer.is_empty() && last_flush.elapsed().as_millis() as u64 >= interval_ms {
                        if !self.flush().await {
                            return;
                        }
                        last_flush = Instant::now();
                    }
                }
            }
        }
    }

    /// バッチサイズのみの受信ループ（フラッシュ間隔なし）
    async fn run_batch_only(&mut self) {
        loop {
            match self.rx.recv().await {
                Ok(event) => {
                    self.buffer.push(Record::from(&event));

                    if self.buffer.len() >= self.config.batch_size && !self.flush().await {
                        return;
                    }
                }
                Err(broadcast::error::RecvError::Closed) => {
                    tracing::debug!("LearningEvent channel closed");
                    return;
                }
                Err(broadcast::error::RecvError::Lagged(n)) => {
                    tracing::warn!(lagged = n, "LearningEventSubscriber lagged behind");
                }
            }
        }
    }

    /// バッファを送信
    ///
    /// Returns: 送信成功なら true、チャンネル閉鎖なら false
    async fn flush(&mut self) -> bool {
        if self.buffer.is_empty() {
            return true;
        }

        let records = std::mem::take(&mut self.buffer);
        let count = records.len();

        match self.record_tx.send(records).await {
            Ok(()) => {
                tracing::debug!(count, "Flushed LearningEvent records to LearningDaemon");
                true
            }
            Err(_) => {
                tracing::warn!("LearningDaemon channel closed");
                false
            }
        }
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::events::{ActionEventBuilder, ActionEventResult, LearningEvent};
    use crate::types::WorkerId;
    use std::time::Duration;

    fn make_action_event(tick: u64, action: &str) -> ActionEvent {
        ActionEventBuilder::new(tick, WorkerId(0), action)
            .result(ActionEventResult::success())
            .duration(Duration::from_millis(10))
            .build()
    }

    fn make_learning_event(model: &str) -> LearningEvent {
        LearningEvent::dependency_graph_inference(model)
            .prompt("test prompt")
            .response("test response")
            .discover_order(vec!["A".into(), "B".into()])
            .success()
            .build()
    }

    // ========================================================================
    // ActionEventSubscriber Tests
    // ========================================================================

    #[tokio::test]
    async fn test_action_subscriber_batch() {
        let (tx, rx) = broadcast::channel::<ActionEvent>(16);
        let (record_tx, mut record_rx) = mpsc::channel::<Vec<Record>>(16);

        let config = EventSubscriberConfig::new()
            .batch_size(3)
            .no_flush_interval();

        let subscriber = ActionEventSubscriber::with_config(rx, record_tx, config);

        let handle = tokio::spawn(async move {
            subscriber.run().await;
        });

        // Send 5 events (should trigger 1 batch of 3)
        for i in 0..5 {
            tx.send(make_action_event(i, &format!("Action{}", i)))
                .unwrap();
        }

        // Wait for batch
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Should receive batch of 3
        let batch = record_rx.try_recv().unwrap();
        assert_eq!(batch.len(), 3);

        // Drop sender to close channel
        drop(tx);

        // Wait for subscriber to finish
        let _ = handle.await;

        // Should receive remaining 2
        let batch = record_rx.try_recv().unwrap();
        assert_eq!(batch.len(), 2);
    }

    #[tokio::test]
    async fn test_action_subscriber_flush_interval() {
        let (tx, rx) = broadcast::channel::<ActionEvent>(16);
        let (record_tx, mut record_rx) = mpsc::channel::<Vec<Record>>(16);

        let config = EventSubscriberConfig::new()
            .batch_size(100) // Large batch size
            .flush_interval_ms(50); // Short flush interval

        let subscriber = ActionEventSubscriber::with_config(rx, record_tx, config);

        let handle = tokio::spawn(async move {
            subscriber.run().await;
        });

        // Send 2 events (won't reach batch size)
        tx.send(make_action_event(0, "Action0")).unwrap();
        tx.send(make_action_event(1, "Action1")).unwrap();

        // Wait for flush interval
        tokio::time::sleep(Duration::from_millis(100)).await;

        // Should receive batch due to flush interval
        let batch = record_rx.try_recv().unwrap();
        assert_eq!(batch.len(), 2);

        drop(tx);
        let _ = handle.await;
    }

    #[tokio::test]
    async fn test_action_subscriber_channel_closed() {
        let (tx, rx) = broadcast::channel::<ActionEvent>(16);
        let (record_tx, record_rx) = mpsc::channel::<Vec<Record>>(16);

        let config = EventSubscriberConfig::new()
            .batch_size(100)
            .no_flush_interval();

        let subscriber = ActionEventSubscriber::with_config(rx, record_tx, config);

        let handle = tokio::spawn(async move {
            subscriber.run().await;
        });

        // Send event
        tx.send(make_action_event(0, "Action0")).unwrap();

        // Drop receiver to close channel
        drop(record_rx);

        // Send more events
        tx.send(make_action_event(1, "Action1")).unwrap();

        // Subscriber should detect closed channel and stop
        tokio::time::sleep(Duration::from_millis(50)).await;

        drop(tx);
        let _ = handle.await;
    }

    // ========================================================================
    // LearningEventSubscriber Tests
    // ========================================================================

    #[tokio::test]
    async fn test_learning_subscriber_batch() {
        let (tx, rx) = broadcast::channel::<LearningEvent>(16);
        let (record_tx, mut record_rx) = mpsc::channel::<Vec<Record>>(16);

        let config = EventSubscriberConfig::new()
            .batch_size(2)
            .no_flush_interval();

        let subscriber = LearningEventSubscriber::with_config(rx, record_tx, config);

        let handle = tokio::spawn(async move {
            subscriber.run().await;
        });

        // Send 3 events (should trigger 1 batch of 2)
        for i in 0..3 {
            tx.send(make_learning_event(&format!("model{}", i)))
                .unwrap();
        }

        // Wait for batch
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Should receive batch of 2
        let batch = record_rx.try_recv().unwrap();
        assert_eq!(batch.len(), 2);

        // All records should be DependencyGraph
        for record in &batch {
            assert!(record.is_dependency_graph());
        }

        // Drop sender to close channel
        drop(tx);

        // Wait for subscriber to finish
        let _ = handle.await;

        // Should receive remaining 1
        let batch = record_rx.try_recv().unwrap();
        assert_eq!(batch.len(), 1);
    }

    #[tokio::test]
    async fn test_learning_subscriber_flush_interval() {
        let (tx, rx) = broadcast::channel::<LearningEvent>(16);
        let (record_tx, mut record_rx) = mpsc::channel::<Vec<Record>>(16);

        let config = EventSubscriberConfig::new()
            .batch_size(100) // Large batch size
            .flush_interval_ms(50); // Short flush interval

        let subscriber = LearningEventSubscriber::with_config(rx, record_tx, config);

        let handle = tokio::spawn(async move {
            subscriber.run().await;
        });

        // Send 1 event (won't reach batch size)
        tx.send(make_learning_event("model")).unwrap();

        // Wait for flush interval
        tokio::time::sleep(Duration::from_millis(100)).await;

        // Should receive batch due to flush interval
        let batch = record_rx.try_recv().unwrap();
        assert_eq!(batch.len(), 1);
        assert!(batch[0].is_dependency_graph());

        drop(tx);
        let _ = handle.await;
    }

    #[tokio::test]
    async fn test_learning_subscriber_converts_to_dependency_graph_record() {
        let (tx, rx) = broadcast::channel::<LearningEvent>(16);
        let (record_tx, mut record_rx) = mpsc::channel::<Vec<Record>>(16);

        let config = EventSubscriberConfig::new()
            .batch_size(1)
            .no_flush_interval();

        let subscriber = LearningEventSubscriber::with_config(rx, record_tx, config);

        let handle = tokio::spawn(async move {
            subscriber.run().await;
        });

        // Send DependencyGraphInference event
        tx.send(make_learning_event("test-model")).unwrap();

        // Wait for batch
        tokio::time::sleep(Duration::from_millis(50)).await;

        let batch = record_rx.try_recv().unwrap();
        assert_eq!(batch.len(), 1);

        // Verify it's a DependencyGraphRecord
        let record = batch[0].as_dependency_graph().unwrap();
        assert_eq!(record.model, "test-model");
        assert_eq!(record.prompt, "test prompt");
        assert_eq!(record.discover_order, vec!["A", "B"]);

        drop(tx);
        let _ = handle.await;
    }

    // ========================================================================
    // Config Tests
    // ========================================================================

    #[tokio::test]
    async fn test_subscriber_config() {
        let config = EventSubscriberConfig::new()
            .batch_size(50)
            .flush_interval_ms(500);

        assert_eq!(config.batch_size, 50);
        assert_eq!(config.flush_interval_ms, Some(500));

        let config2 = EventSubscriberConfig::new().no_flush_interval();
        assert_eq!(config2.flush_interval_ms, None);
    }
}