drasi_lib/reactions/common/

base.rs

// Copyright 2025 The Drasi Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Base implementation for common reaction functionality.
//!
//! This module provides `ReactionBase` which encapsulates common patterns
//! used across all reaction implementations:
//! - Query subscription management
//! - Priority queue handling
//! - Task lifecycle management
//! - Component status tracking
//! - Event reporting
//!
//! # Plugin Architecture
//!
//! ReactionBase is designed to be used by reaction plugins. Each plugin:
//! 1. Defines its own typed configuration struct
//! 2. Creates a ReactionBase with ReactionBaseParams
//! 3. Implements the Reaction trait delegating to ReactionBase methods

use anyhow::Result;
use log::{debug, error, info, warn};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::Instrument;

use crate::channels::priority_queue::PriorityQueue;
use crate::channels::{ComponentStatus, QueryResult};
use crate::component_graph::ComponentStatusHandle;
use crate::context::ReactionRuntimeContext;
use crate::identity::IdentityProvider;
use crate::state_store::StateStoreProvider;

/// Parameters for creating a ReactionBase instance.
///
/// This struct contains only the information that ReactionBase needs to function.
/// Plugin-specific configuration should remain in the plugin crate.
///
/// # Example
///
/// ```ignore
/// use drasi_lib::reactions::common::base::{ReactionBase, ReactionBaseParams};
///
/// let params = ReactionBaseParams::new("my-reaction", vec!["query1".to_string()])
///     .with_priority_queue_capacity(5000)
///     .with_auto_start(true);
///
/// let base = ReactionBase::new(params);
/// ```
#[derive(Debug, Clone)]
pub struct ReactionBaseParams {
    /// Unique identifier for the reaction
    pub id: String,
    /// List of query IDs this reaction subscribes to
    pub queries: Vec<String>,
    /// Priority queue capacity - defaults to 10000
    pub priority_queue_capacity: Option<usize>,
    /// Whether this reaction should auto-start - defaults to true
    pub auto_start: bool,
}

impl ReactionBaseParams {
    /// Create new params with ID and queries, using defaults for everything else
    pub fn new(id: impl Into<String>, queries: Vec<String>) -> Self {
        Self {
            id: id.into(),
            queries,
            priority_queue_capacity: None,
            auto_start: true, // Default to true like queries
        }
    }

    /// Set the priority queue capacity
    pub fn with_priority_queue_capacity(mut self, capacity: usize) -> Self {
        self.priority_queue_capacity = Some(capacity);
        self
    }

    /// Set whether this reaction should auto-start
    pub fn with_auto_start(mut self, auto_start: bool) -> Self {
        self.auto_start = auto_start;
        self
    }
}

/// Base implementation for common reaction functionality
pub struct ReactionBase {
    /// Reaction identifier
    pub id: String,
    /// List of query IDs to subscribe to
    pub queries: Vec<String>,
    /// Whether this reaction should auto-start
    pub auto_start: bool,
    /// Component status handle — always available, wired to graph during initialize().
    status_handle: ComponentStatusHandle,
    /// Runtime context (set by initialize())
    context: Arc<RwLock<Option<ReactionRuntimeContext>>>,
    /// State store provider (extracted from context for convenience)
    state_store: Arc<RwLock<Option<Arc<dyn StateStoreProvider>>>>,
    /// Priority queue for timestamp-ordered result processing
    pub priority_queue: PriorityQueue<QueryResult>,
    /// Handles to subscription forwarder tasks
    pub subscription_tasks: Arc<RwLock<Vec<tokio::task::JoinHandle<()>>>>,
    /// Handle to the main processing task
    pub processing_task: Arc<RwLock<Option<tokio::task::JoinHandle<()>>>>,
    /// Sender for shutdown signal to processing task
    pub shutdown_tx: Arc<RwLock<Option<tokio::sync::oneshot::Sender<()>>>>,
    /// Optional identity provider for credential management.
    /// Set either programmatically (via `set_identity_provider`) or automatically
    /// from the runtime context during `initialize()`.
    identity_provider: Arc<RwLock<Option<Arc<dyn IdentityProvider>>>>,
}

impl ReactionBase {
    /// Create a new ReactionBase with the given parameters
    ///
    /// Dependencies (query subscriber, state store, graph) are not required during
    /// construction - they will be provided via `initialize()` when the reaction is added to DrasiLib.
    pub fn new(params: ReactionBaseParams) -> Self {
        Self {
            priority_queue: PriorityQueue::new(params.priority_queue_capacity.unwrap_or(10000)),
            id: params.id.clone(),
            queries: params.queries,
            auto_start: params.auto_start,
            status_handle: ComponentStatusHandle::new(&params.id),
            context: Arc::new(RwLock::new(None)), // Set by initialize()
            state_store: Arc::new(RwLock::new(None)), // Extracted from context
            subscription_tasks: Arc::new(RwLock::new(Vec::new())),
            processing_task: Arc::new(RwLock::new(None)),
            shutdown_tx: Arc::new(RwLock::new(None)),
            identity_provider: Arc::new(RwLock::new(None)),
        }
    }

    /// Initialize the reaction with runtime context.
    ///
    /// This method is called automatically by DrasiLib's `add_reaction()` method.
    /// Plugin developers do not need to call this directly.
    ///
    /// The context provides access to:
    /// - `reaction_id`: The reaction's unique identifier
    /// - `state_store`: Optional persistent state storage
    /// - `update_tx`: mpsc sender for fire-and-forget status updates to the graph
    pub async fn initialize(&self, context: ReactionRuntimeContext) {
        // Store context for later use
        *self.context.write().await = Some(context.clone());

        // Wire the status handle to the graph update channel
        self.status_handle.wire(context.update_tx.clone()).await;

        if let Some(state_store) = context.state_store.as_ref() {
            *self.state_store.write().await = Some(state_store.clone());
        }

        // Store identity provider from context if not already set programmatically
        if let Some(ip) = context.identity_provider.as_ref() {
            let mut guard = self.identity_provider.write().await;
            if guard.is_none() {
                *guard = Some(ip.clone());
            }
        }
    }

    /// Get the runtime context if initialized.
    ///
    /// Returns `None` if `initialize()` has not been called yet.
    pub async fn context(&self) -> Option<ReactionRuntimeContext> {
        self.context.read().await.clone()
    }

    /// Get the state store if configured.
    ///
    /// Returns `None` if no state store was provided in the context.
    pub async fn state_store(&self) -> Option<Arc<dyn StateStoreProvider>> {
        self.state_store.read().await.clone()
    }

    /// Get the identity provider if set.
    ///
    /// Returns the identity provider set either programmatically via
    /// `set_identity_provider()` or from the runtime context during `initialize()`.
    /// Programmatically-set providers take precedence over context providers.
    pub async fn identity_provider(&self) -> Option<Arc<dyn IdentityProvider>> {
        self.identity_provider.read().await.clone()
    }

    /// Set the identity provider programmatically.
    ///
    /// This is typically called during reaction construction when the provider
    /// is available from configuration (e.g., `with_identity_provider()` builder).
    /// Providers set this way take precedence over context-injected providers.
    pub async fn set_identity_provider(&self, provider: Arc<dyn IdentityProvider>) {
        *self.identity_provider.write().await = Some(provider);
    }

    /// Get whether this reaction should auto-start
    pub fn get_auto_start(&self) -> bool {
        self.auto_start
    }

    /// Clone the ReactionBase with shared Arc references
    ///
    /// This creates a new ReactionBase that shares the same underlying
    /// data through Arc references. Useful for passing to spawned tasks.
    pub fn clone_shared(&self) -> Self {
        Self {
            id: self.id.clone(),
            queries: self.queries.clone(),
            auto_start: self.auto_start,
            status_handle: self.status_handle.clone(),
            context: self.context.clone(),
            state_store: self.state_store.clone(),
            priority_queue: self.priority_queue.clone(),
            subscription_tasks: self.subscription_tasks.clone(),
            processing_task: self.processing_task.clone(),
            shutdown_tx: self.shutdown_tx.clone(),
            identity_provider: self.identity_provider.clone(),
        }
    }

    /// Create a shutdown channel and store the sender
    ///
    /// Returns the receiver which should be passed to the processing task.
    /// The sender is stored internally and will be triggered by `stop_common()`.
    ///
    /// This should be called before spawning the processing task.
    pub async fn create_shutdown_channel(&self) -> tokio::sync::oneshot::Receiver<()> {
        let (tx, rx) = tokio::sync::oneshot::channel();
        *self.shutdown_tx.write().await = Some(tx);
        rx
    }

    /// Get the reaction ID
    pub fn get_id(&self) -> &str {
        &self.id
    }

    /// Get the query IDs
    pub fn get_queries(&self) -> &[String] {
        &self.queries
    }

    /// Get current status.
    pub async fn get_status(&self) -> ComponentStatus {
        self.status_handle.get_status().await
    }

    /// Returns a clonable [`ComponentStatusHandle`] for use in spawned tasks.
    ///
    /// The handle can both read and write the component's status and automatically
    /// notifies the graph on every status change (after `initialize()`).
    pub fn status_handle(&self) -> ComponentStatusHandle {
        self.status_handle.clone()
    }

    /// Set the component's status — updates local state AND notifies the graph.
    ///
    /// This is the single canonical way to change a reaction's status.
    pub async fn set_status(&self, status: ComponentStatus, message: Option<String>) {
        self.status_handle.set_status(status, message).await;
    }

    /// Enqueue a query result for processing.
    ///
    /// The host calls this to forward query results to the reaction's priority queue.
    /// Results are processed in timestamp order by the reaction's processing task.
    pub async fn enqueue_query_result(&self, result: QueryResult) -> anyhow::Result<()> {
        self.priority_queue.enqueue_wait(Arc::new(result)).await;
        Ok(())
    }

    /// Perform common cleanup operations
    ///
    /// This method handles:
    /// 1. Sending shutdown signal to processing task (for graceful termination)
    /// 2. Aborting all subscription forwarder tasks
    /// 3. Waiting for or aborting the processing task
    /// 4. Draining the priority queue
    pub async fn stop_common(&self) -> Result<()> {
        info!("Stopping reaction: {}", self.id);

        // Send shutdown signal to processing task (if it's using tokio::select!)
        if let Some(tx) = self.shutdown_tx.write().await.take() {
            let _ = tx.send(());
        }

        // Abort all subscription forwarder tasks
        let mut subscription_tasks = self.subscription_tasks.write().await;
        for task in subscription_tasks.drain(..) {
            task.abort();
        }
        drop(subscription_tasks);

        // Wait for the processing task to complete (with timeout), or abort it
        let mut processing_task = self.processing_task.write().await;
        if let Some(mut task) = processing_task.take() {
            // Give the task a short time to respond to the shutdown signal
            match tokio::time::timeout(std::time::Duration::from_secs(2), &mut task).await {
                Ok(Ok(())) => {
                    debug!("[{}] Processing task completed gracefully", self.id);
                }
                Ok(Err(e)) => {
                    // Task was aborted or panicked
                    debug!("[{}] Processing task ended: {}", self.id, e);
                }
                Err(_) => {
                    // Timeout - task didn't respond to shutdown signal
                    warn!(
                        "[{}] Processing task did not respond to shutdown signal within timeout, aborting",
                        self.id
                    );
                    task.abort();
                }
            }
        }
        drop(processing_task);

        // Drain the priority queue
        let drained_events = self.priority_queue.drain().await;
        if !drained_events.is_empty() {
            info!(
                "[{}] Drained {} pending events from priority queue",
                self.id,
                drained_events.len()
            );
        }

        self.set_status(
            ComponentStatus::Stopped,
            Some(format!("Reaction '{}' stopped", self.id)),
        )
        .await;
        info!("Reaction '{}' stopped", self.id);

        Ok(())
    }

    /// Clear the reaction's state store partition.
    ///
    /// This is called during deprovision to remove all persisted state
    /// associated with this reaction. Reactions that override `deprovision()`
    /// can call this to clean up their state store.
    pub async fn deprovision_common(&self) -> Result<()> {
        info!("Deprovisioning reaction '{}'", self.id);
        if let Some(store) = self.state_store().await {
            let count = store.clear_store(&self.id).await.map_err(|e| {
                anyhow::anyhow!(
                    "Failed to clear state store for reaction '{}': {}",
                    self.id,
                    e
                )
            })?;
            info!(
                "Cleared {} keys from state store for reaction '{}'",
                count, self.id
            );
        }
        Ok(())
    }

    /// Set the processing task handle
    pub async fn set_processing_task(&self, task: tokio::task::JoinHandle<()>) {
        *self.processing_task.write().await = Some(task);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::time::Duration;
    use tokio::sync::mpsc;

    #[tokio::test]
    async fn test_reaction_base_creation() {
        let params = ReactionBaseParams::new("test-reaction", vec!["query1".to_string()])
            .with_priority_queue_capacity(5000);

        let base = ReactionBase::new(params);
        assert_eq!(base.id, "test-reaction");
        assert_eq!(base.get_status().await, ComponentStatus::Stopped);
    }

    #[tokio::test]
    async fn test_status_transitions() {
        use crate::context::ReactionRuntimeContext;

        let (graph, _rx) = crate::component_graph::ComponentGraph::new("test-instance");
        let update_tx = graph.update_sender();
        let graph = Arc::new(RwLock::new(graph));
        let params = ReactionBaseParams::new("test-reaction", vec![]);

        let base = ReactionBase::new(params);

        // Create context and initialize
        let context =
            ReactionRuntimeContext::new("test-instance", "test-reaction", None, update_tx, None);
        base.initialize(context).await;

        // Test status transition
        base.set_status(ComponentStatus::Starting, Some("Starting test".to_string()))
            .await;

        assert_eq!(base.get_status().await, ComponentStatus::Starting);

        // Check event was sent via graph broadcast
        let mut event_rx = graph.read().await.subscribe();
        // The status was already set; emit another event to verify the graph path works
        base.set_status(ComponentStatus::Running, Some("Running test".to_string()))
            .await;

        assert_eq!(base.get_status().await, ComponentStatus::Running);
    }

    #[tokio::test]
    async fn test_priority_queue_operations() {
        let params =
            ReactionBaseParams::new("test-reaction", vec![]).with_priority_queue_capacity(10);

        let base = ReactionBase::new(params);

        // Create a test query result
        let query_result = QueryResult::new(
            "test-query".to_string(),
            chrono::Utc::now(),
            vec![],
            Default::default(),
        );

        // Enqueue result
        let enqueued = base.priority_queue.enqueue(Arc::new(query_result)).await;
        assert!(enqueued);

        // Drain queue
        let drained = base.priority_queue.drain().await;
        assert_eq!(drained.len(), 1);
    }

    #[tokio::test]
    async fn test_event_without_initialization() {
        // Test that set_status works even without context initialization
        let params = ReactionBaseParams::new("test-reaction", vec![]);

        let base = ReactionBase::new(params);

        // This should succeed without panicking (silently updates local only when handle is None)
        base.set_status(ComponentStatus::Starting, None).await;
    }

    // =============================================================================
    // Shutdown Channel Tests
    // =============================================================================

    #[tokio::test]
    async fn test_create_shutdown_channel() {
        let params = ReactionBaseParams::new("test-reaction", vec![]);
        let base = ReactionBase::new(params);

        // Initially no shutdown_tx
        assert!(base.shutdown_tx.read().await.is_none());

        // Create channel
        let rx = base.create_shutdown_channel().await;

        // Verify tx is stored
        assert!(base.shutdown_tx.read().await.is_some());

        // Verify receiver is valid (dropping it should not panic)
        drop(rx);
    }

    #[tokio::test]
    async fn test_shutdown_channel_signal() {
        let params = ReactionBaseParams::new("test-reaction", vec![]);
        let base = ReactionBase::new(params);

        let mut rx = base.create_shutdown_channel().await;

        // Send signal
        if let Some(tx) = base.shutdown_tx.write().await.take() {
            tx.send(()).unwrap();
        }

        // Verify signal received
        let result = rx.try_recv();
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_shutdown_channel_replaced_on_second_create() {
        let params = ReactionBaseParams::new("test-reaction", vec![]);
        let base = ReactionBase::new(params);

        // Create first channel
        let _rx1 = base.create_shutdown_channel().await;

        // Create second channel (should replace the first)
        let mut rx2 = base.create_shutdown_channel().await;

        // Send signal - should go to second channel
        if let Some(tx) = base.shutdown_tx.write().await.take() {
            tx.send(()).unwrap();
        }

        // Second receiver should get the signal
        let result = rx2.try_recv();
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_stop_common_sends_shutdown_signal() {
        let params = ReactionBaseParams::new("test-reaction", vec![]);
        let base = ReactionBase::new(params);

        let mut rx = base.create_shutdown_channel().await;

        // Spawn a task that waits for shutdown
        let shutdown_received = Arc::new(AtomicBool::new(false));
        let shutdown_flag = shutdown_received.clone();

        let task = tokio::spawn(async move {
            tokio::select! {
                _ = &mut rx => {
                    shutdown_flag.store(true, Ordering::SeqCst);
                }
            }
        });

        base.set_processing_task(task).await;

        // Call stop_common - should send shutdown signal and await the task
        let _ = base.stop_common().await;

        // stop_common awaits the processing task, so the flag should already be set
        assert!(
            shutdown_received.load(Ordering::SeqCst),
            "Processing task should have received shutdown signal"
        );
    }

    #[tokio::test]
    async fn test_graceful_shutdown_timing() {
        let params = ReactionBaseParams::new("test-reaction", vec![]);
        let base = ReactionBase::new(params);

        let rx = base.create_shutdown_channel().await;

        // Spawn task that uses select! pattern like real reactions
        let task = tokio::spawn(async move {
            let mut shutdown_rx = rx;
            loop {
                tokio::select! {
                    biased;
                    _ = &mut shutdown_rx => {
                        break;
                    }
                    _ = tokio::time::sleep(Duration::from_secs(10)) => {
                        // Simulates waiting on priority_queue.dequeue()
                    }
                }
            }
        });

        base.set_processing_task(task).await;

        // Measure shutdown time
        let start = std::time::Instant::now();
        let _ = base.stop_common().await;
        let elapsed = start.elapsed();

        // Should complete quickly (< 500ms), not hit 2s timeout
        assert!(
            elapsed < Duration::from_millis(500),
            "Shutdown took {elapsed:?}, expected < 500ms. Task may not be responding to shutdown signal."
        );
    }

    #[tokio::test]
    async fn test_stop_common_without_shutdown_channel() {
        // Test that stop_common works even if no shutdown channel was created
        let params = ReactionBaseParams::new("test-reaction", vec![]);
        let base = ReactionBase::new(params);

        // Don't create shutdown channel - just spawn a short-lived task
        let task = tokio::spawn(async {
            tokio::time::sleep(Duration::from_millis(10)).await;
        });

        base.set_processing_task(task).await;

        // stop_common should still work
        let result = base.stop_common().await;
        assert!(result.is_ok());
    }

    // =============================================================================
    // Accessor Tests
    // =============================================================================

    #[tokio::test]
    async fn test_get_id() {
        let params = ReactionBaseParams::new("my-reaction-42", vec![]);
        let base = ReactionBase::new(params);
        assert_eq!(base.get_id(), "my-reaction-42");
    }

    #[tokio::test]
    async fn test_get_queries() {
        let queries = vec!["query-a".to_string(), "query-b".to_string(), "query-c".to_string()];
        let params = ReactionBaseParams::new("r1", queries.clone());
        let base = ReactionBase::new(params);
        assert_eq!(base.get_queries(), &queries[..]);
    }

    #[tokio::test]
    async fn test_get_queries_empty() {
        let params = ReactionBaseParams::new("r1", vec![]);
        let base = ReactionBase::new(params);
        assert!(base.get_queries().is_empty());
    }

    #[tokio::test]
    async fn test_get_auto_start_default_true() {
        let params = ReactionBaseParams::new("r1", vec![]);
        let base = ReactionBase::new(params);
        assert!(base.get_auto_start());
    }

    #[tokio::test]
    async fn test_get_auto_start_override_false() {
        let params = ReactionBaseParams::new("r1", vec![]).with_auto_start(false);
        let base = ReactionBase::new(params);
        assert!(!base.get_auto_start());
    }

    // =============================================================================
    // Context / State Store / Identity Provider Tests
    // =============================================================================

    #[tokio::test]
    async fn test_context_none_before_initialize() {
        let params = ReactionBaseParams::new("r1", vec![]);
        let base = ReactionBase::new(params);
        assert!(base.context().await.is_none());
    }

    #[tokio::test]
    async fn test_context_some_after_initialize() {
        let (graph, _rx) = crate::component_graph::ComponentGraph::new("inst");
        let update_tx = graph.update_sender();
        let context = ReactionRuntimeContext::new("inst", "r1", None, update_tx, None);

        let params = ReactionBaseParams::new("r1", vec![]);
        let base = ReactionBase::new(params);
        base.initialize(context).await;

        let ctx = base.context().await;
        assert!(ctx.is_some());
        assert_eq!(ctx.unwrap().reaction_id, "r1");
    }

    #[tokio::test]
    async fn test_state_store_none_when_not_configured() {
        let params = ReactionBaseParams::new("r1", vec![]);
        let base = ReactionBase::new(params);
        assert!(base.state_store().await.is_none());
    }

    #[tokio::test]
    async fn test_state_store_none_after_initialize_without_store() {
        let (graph, _rx) = crate::component_graph::ComponentGraph::new("inst");
        let update_tx = graph.update_sender();
        let context = ReactionRuntimeContext::new("inst", "r1", None, update_tx, None);

        let params = ReactionBaseParams::new("r1", vec![]);
        let base = ReactionBase::new(params);
        base.initialize(context).await;

        assert!(base.state_store().await.is_none());
    }

    #[tokio::test]
    async fn test_identity_provider_none_by_default() {
        let params = ReactionBaseParams::new("r1", vec![]);
        let base = ReactionBase::new(params);
        assert!(base.identity_provider().await.is_none());
    }

    // =============================================================================
    // Status Handle Tests
    // =============================================================================

    #[tokio::test]
    async fn test_status_handle_returns_handle() {
        let params = ReactionBaseParams::new("r1", vec![]);
        let base = ReactionBase::new(params);

        let handle = base.status_handle();
        // The handle should share the same status as the base
        assert_eq!(handle.get_status().await, ComponentStatus::Stopped);

        // Mutating via handle should be visible via base
        handle.set_status(ComponentStatus::Running, None).await;
        assert_eq!(base.get_status().await, ComponentStatus::Running);
    }

    // =============================================================================
    // Deprovision Tests
    // =============================================================================

    #[tokio::test]
    async fn test_deprovision_common_noop_without_state_store() {
        let params = ReactionBaseParams::new("r1", vec![]);
        let base = ReactionBase::new(params);
        // Should succeed without panicking when no state store is configured
        let result = base.deprovision_common().await;
        assert!(result.is_ok());
    }

    // =============================================================================
    // Processing Task Tests
    // =============================================================================

    #[tokio::test]
    async fn test_set_processing_task_stores_handle() {
        let params = ReactionBaseParams::new("r1", vec![]);
        let base = ReactionBase::new(params);

        // Initially no processing task
        assert!(base.processing_task.read().await.is_none());

        let task = tokio::spawn(async {
            tokio::time::sleep(Duration::from_secs(60)).await;
        });

        base.set_processing_task(task).await;

        // Now it should be stored
        assert!(base.processing_task.read().await.is_some());

        // Clean up: abort the long-running task
        let task = base.processing_task.write().await.take();
        if let Some(t) = task {
            t.abort();
        }
    }
}