pollen_rs/

builder.rs

//! Builder API for Pollen scheduler.

use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use pollen_claimer::{Claimer, ClaimerConfig};
use pollen_clock::SharedClock;
use pollen_executor::{Executor, TokioExecutor};
use pollen_membership::{ChitchatMembership, Membership};
use pollen_router::{ConsistentHashRouter, TaskRouter};
use pollen_scheduler::{DefaultScheduler, SharedScheduler};
use pollen_store::{MemoryStore, SqliteStore, StorageConfig, StoreBackend};
use pollen_types::{NodeId, Result, Schedule, TaskDef};
use tokio_util::sync::CancellationToken;
use tracing::info;

/// Storage configuration
#[derive(Debug, Clone)]
pub enum Storage {
    /// SQLite file-based storage
    Sqlite(PathBuf),
    /// In-memory storage (for testing)
    Memory,
}

impl Storage {
    /// Create SQLite storage configuration
    pub fn sqlite<P: Into<PathBuf>>(path: P) -> Self {
        Storage::Sqlite(path.into())
    }

    /// Create in-memory storage configuration
    pub fn memory() -> Self {
        Storage::Memory
    }
}

/// Cluster configuration for the builder
#[derive(Debug, Clone)]
pub struct ClusterConfig {
    /// Local bind address for cluster communication
    pub bind_addr: SocketAddr,
    /// Seed nodes to join the cluster
    pub seeds: Vec<SocketAddr>,
    /// Cluster name for isolation
    pub cluster_name: String,
    /// Node metadata to propagate via gossip
    pub metadata: Vec<(String, String)>,
}

impl Default for ClusterConfig {
    fn default() -> Self {
        Self {
            bind_addr: "0.0.0.0:7000".parse().unwrap(),
            seeds: Vec::new(),
            cluster_name: "pollen".to_string(),
            metadata: Vec::new(),
        }
    }
}

/// Executor configuration
#[derive(Debug, Clone)]
pub struct ExecutorSettings {
    /// Maximum concurrent task executions
    pub max_concurrent: usize,
    /// Default timeout for task execution
    pub default_timeout: Duration,
}

impl Default for ExecutorSettings {
    fn default() -> Self {
        Self {
            max_concurrent: 100,
            default_timeout: Duration::from_secs(300),
        }
    }
}

/// Builder for constructing a Pollen scheduler instance
pub struct PollenBuilder {
    storage: Storage,
    cluster: Option<ClusterConfig>,
    executor: ExecutorSettings,
}

impl PollenBuilder {
    /// Create a new builder with default configuration
    pub fn new() -> Self {
        Self {
            storage: Storage::Memory,
            cluster: None,
            executor: ExecutorSettings::default(),
        }
    }

    /// Configure storage backend
    pub fn storage(mut self, storage: Storage) -> Self {
        self.storage = storage;
        self
    }

    /// Configure cluster mode
    pub fn cluster(mut self, config: ClusterConfig) -> Self {
        self.cluster = Some(config);
        self
    }

    /// Configure executor
    pub fn executor(mut self, config: ExecutorSettings) -> Self {
        self.executor = config;
        self
    }

    /// Build the Pollen scheduler instance
    pub async fn build(self) -> Result<Pollen> {
        // Generate node ID
        let node_id = NodeId::new();
        info!("Starting Pollen with node ID: {:?}", node_id);

        // 1. Initialize foundation layer

        // Create HLC clock
        let clock = pollen_clock::new_clock_with_id(node_id);

        // Create storage backend
        let store: Arc<StoreBackend> = match &self.storage {
            Storage::Sqlite(path) => {
                let config = StorageConfig {
                    path: path.to_string_lossy().to_string(),
                    ..Default::default()
                };
                let sqlite = SqliteStore::open(&config)?;
                Arc::new(StoreBackend::Sqlite(sqlite))
            }
            Storage::Memory => Arc::new(StoreBackend::Memory(MemoryStore::new())),
        };

        // Create executor
        let executor_config = pollen_executor::ExecutorConfig {
            max_concurrent: self.executor.max_concurrent,
            default_timeout: self.executor.default_timeout,
        };
        let executor: Arc<dyn Executor> = Arc::new(TokioExecutor::new(executor_config));

        // 2. Initialize cluster layer (if configured)
        #[allow(clippy::type_complexity)]
        let (membership, router): (
            Option<Arc<dyn Membership>>,
            Option<Arc<dyn TaskRouter>>,
        ) = if let Some(ref cluster) = self.cluster {
            // Convert builder ClusterConfig to pollen_types::ClusterConfig
            let types_config = pollen_types::ClusterConfig {
                bind_addr: cluster.bind_addr,
                seeds: cluster.seeds.clone(),
                cluster_name: cluster.cluster_name.clone(),
                metadata: cluster.metadata.iter().cloned().collect(),
                ..Default::default()
            };

            // Create chitchat-based membership
            let membership: Arc<dyn Membership> =
                Arc::new(ChitchatMembership::new(types_config, node_id).await?);

            // Create consistent hash router
            let router: Arc<dyn TaskRouter> =
                Arc::new(ConsistentHashRouter::new(node_id, Arc::clone(&membership)));

            info!(
                "Pollen cluster mode initialized at {}",
                cluster.bind_addr
            );

            (Some(membership), Some(router))
        } else {
            (None, None)
        };

        // 3. Initialize coordination layer

        // Create scheduler
        let scheduler = Arc::new(DefaultScheduler::new(
            clock.clone(),
            Arc::clone(&store),
            None, // CRDT store (future work)
        ));

        // Start scheduler background tasks (instance generation)
        scheduler.clone().start();

        let scheduler: SharedScheduler = scheduler;

        // Create claimer
        let claimer_config = ClaimerConfig::default();
        let claimer = Arc::new(Claimer::new(
            node_id,
            Arc::clone(&store),
            router.clone(),
            Arc::clone(&executor),
            Arc::clone(&scheduler),
            membership.clone(),
            claimer_config,
        ));

        // Create shutdown token
        let shutdown = CancellationToken::new();

        // 4. Start background tasks
        let claimer_handle = claimer.clone();
        let shutdown_clone = shutdown.clone();
        tokio::spawn(async move {
            claimer_handle.run(shutdown_clone).await;
        });

        let mode = if self.cluster.is_some() {
            "cluster"
        } else {
            "single-node"
        };
        info!("Pollen scheduler initialized in {} mode", mode);

        Ok(Pollen {
            node_id,
            _clock: clock,
            _store: store,
            _executor: executor,
            scheduler,
            claimer,
            membership,
            shutdown,
            _storage_config: self.storage,
            _cluster_config: self.cluster,
        })
    }
}

impl Default for PollenBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// Main Pollen scheduler instance
pub struct Pollen {
    node_id: NodeId,
    _clock: SharedClock,
    _store: Arc<StoreBackend>,
    _executor: Arc<dyn Executor>,
    scheduler: SharedScheduler,
    #[allow(dead_code)] // Used for graceful shutdown and future cluster mode
    claimer: Arc<Claimer>,
    membership: Option<Arc<dyn Membership>>,
    shutdown: CancellationToken,
    _storage_config: Storage,
    _cluster_config: Option<ClusterConfig>,
}

impl Pollen {
    /// Create a new builder
    pub fn builder() -> PollenBuilder {
        PollenBuilder::new()
    }

    /// Get the local node ID
    pub fn node_id(&self) -> NodeId {
        self.node_id
    }

    /// Check if running in cluster mode
    pub fn is_cluster_mode(&self) -> bool {
        self._cluster_config.is_some()
    }

    /// Define a new task
    pub fn define(&self, name: &str) -> TaskBuilder {
        TaskBuilder::new(name.to_string(), Arc::clone(&self.scheduler))
    }

    /// Get a task by name
    pub fn get_task(&self, name: &str) -> Option<TaskDef> {
        self.scheduler.get_task_by_name(name)
    }

    /// Trigger immediate execution of a task
    pub async fn trigger(&self, name: &str) -> Result<()> {
        let task = self
            .scheduler
            .get_task_by_name(name)
            .ok_or_else(|| pollen_types::PollenError::TaskNameNotFound(name.to_string()))?;
        self.scheduler.trigger(&task.id, None).await?;
        Ok(())
    }

    /// Trigger a task with payload
    pub async fn trigger_with_payload(&self, name: &str, payload: bytes::Bytes) -> Result<()> {
        let task = self
            .scheduler
            .get_task_by_name(name)
            .ok_or_else(|| pollen_types::PollenError::TaskNameNotFound(name.to_string()))?;
        self.scheduler.trigger(&task.id, Some(payload)).await?;
        Ok(())
    }

    /// Get task information
    pub fn task_info(&self, name: &str) -> Option<TaskInfo> {
        self.scheduler.get_task_by_name(name).map(|task| TaskInfo {
            id: task.id.clone(),
            name: task.name.clone(),
            enabled: task.enabled,
            last_run: None, // TODO: Query from store
            next_run: None, // TODO: Compute from schedule
        })
    }

    /// List all registered tasks
    pub fn list_tasks(&self) -> Vec<TaskDef> {
        self.scheduler.list_tasks()
    }

    /// Get cluster state (if in cluster mode)
    pub fn cluster_state(&self) -> Option<ClusterState> {
        if let Some(ref membership) = self.membership {
            let members = membership.alive_members();
            Some(ClusterState {
                local_node: self.node_id,
                nodes: members
                    .into_iter()
                    .map(|m| NodeInfo {
                        id: m.id,
                        addr: m.addr,
                        state: match m.state {
                            pollen_types::MemberState::Alive => NodeState::Alive,
                            pollen_types::MemberState::Suspect => NodeState::Suspect,
                            pollen_types::MemberState::Dead | pollen_types::MemberState::Left => {
                                NodeState::Dead
                            }
                        },
                    })
                    .collect(),
            })
        } else {
            None
        }
    }

    /// Get the number of nodes in the cluster
    pub fn cluster_size(&self) -> usize {
        self.membership
            .as_ref()
            .map(|m| m.alive_members().len())
            .unwrap_or(1)
    }

    /// Shutdown the scheduler gracefully
    pub async fn shutdown(self) {
        info!("Shutting down Pollen scheduler...");
        self.shutdown.cancel();

        // Shutdown membership if in cluster mode
        if let Some(membership) = &self.membership {
            membership.shutdown().await;
        }

        // Give background tasks time to complete
        tokio::time::sleep(Duration::from_millis(100)).await;
        info!("Pollen scheduler shutdown complete");
    }
}

/// Task builder for fluent task definition
pub struct TaskBuilder {
    name: String,
    scheduler: SharedScheduler,
    schedule: Option<Schedule>,
    handler: Option<Arc<dyn pollen_executor::TaskHandler>>,
    retry_policy: Option<pollen_types::RetryPolicy>,
    timeout: Option<Duration>,
}

impl TaskBuilder {
    fn new(name: String, scheduler: SharedScheduler) -> Self {
        Self {
            name,
            scheduler,
            schedule: None,
            handler: None,
            retry_policy: None,
            timeout: None,
        }
    }

    /// Set the task schedule
    pub fn schedule(mut self, schedule: Schedule) -> Self {
        self.schedule = Some(schedule);
        self
    }

    /// Set the task handler
    pub fn handler<F, Fut>(mut self, handler: F) -> Self
    where
        F: Fn() -> Fut + Send + Sync + 'static,
        Fut: std::future::Future<Output = Result<()>> + Send + 'static,
    {
        self.handler = Some(pollen_executor::simple_handler(handler));
        self
    }

    /// Set retry policy
    pub fn retry(mut self, policy: pollen_types::RetryPolicy) -> Self {
        self.retry_policy = Some(policy);
        self
    }

    /// Set timeout
    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(timeout);
        self
    }

    /// Register the task
    pub async fn register(self) -> Result<()> {
        let schedule = self.schedule.ok_or_else(|| {
            pollen_types::PollenError::Config("Task schedule is required".to_string())
        })?;

        let handler = self.handler.ok_or_else(|| {
            pollen_types::PollenError::Config("Task handler is required".to_string())
        })?;

        let mut task = TaskDef::new(&self.name, schedule);

        if let Some(retry) = self.retry_policy {
            task = task.with_retry(retry);
        }

        if let Some(timeout) = self.timeout {
            task = task.with_timeout(timeout);
        }

        self.scheduler.register(task, handler).await
    }
}

/// Task information
#[derive(Debug, Clone)]
pub struct TaskInfo {
    pub id: pollen_types::TaskId,
    pub name: String,
    pub enabled: bool,
    pub last_run: Option<chrono::DateTime<chrono::Utc>>,
    pub next_run: Option<chrono::DateTime<chrono::Utc>>,
}

/// Cluster state information
#[derive(Debug, Clone)]
pub struct ClusterState {
    pub local_node: pollen_types::NodeId,
    pub nodes: Vec<NodeInfo>,
}

/// Node information in cluster
#[derive(Debug, Clone)]
pub struct NodeInfo {
    pub id: pollen_types::NodeId,
    pub addr: SocketAddr,
    pub state: NodeState,
}

/// Node state in cluster
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NodeState {
    Alive,
    Suspect,
    Dead,
}