timerwheel 0.1.0

// Copyright © 2026-present The Timerwheel Authors. All rights reserved.
//
// 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.

use std::fmt;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::thread::{self, JoinHandle};

use crate::error::{Error, Result};
use crate::executor::builder::PoolBuilder;
use crate::executor::metrics::{ExecutorMetrics, ExecutorMetricsInner};
use crate::executor::task::BoxTask;
use crate::executor::worker::{self, Message};
use crate::executor::{Executor, RejectedTask};
use crate::policy::RejectPolicy;

/// Bounded worker executor backed by standard library threads.
pub struct Pool {
    sender: mpsc::SyncSender<Message>,
    workers: Mutex<Vec<JoinHandle<()>>>,
    metrics: Arc<ExecutorMetricsInner>,
    reject_policy: RejectPolicy,
    closed: AtomicBool,
}

impl Pool {
    /// Returns a builder for a worker pool.
    pub fn builder() -> PoolBuilder {
        PoolBuilder::default()
    }

    pub(crate) fn new(builder: PoolBuilder) -> Result<Self> {
        let (sender, receiver) = mpsc::sync_channel(builder.queue_capacity);
        // The standard receiver is single-consumer, so workers share it behind
        // a mutex. Task execution happens after recv returns; the lock is not
        // held while user code runs.
        let receiver = Arc::new(Mutex::new(receiver));
        let metrics = Arc::new(ExecutorMetricsInner::new(
            builder.workers,
            Arc::clone(&builder.metric_sink),
        ));
        let mut workers = Vec::with_capacity(builder.workers);

        for index in 0..builder.workers {
            let worker_receiver = Arc::clone(&receiver);
            let worker_metrics = Arc::clone(&metrics);
            let panic_handler = Arc::clone(&builder.panic_handler);
            let handle = thread::Builder::new()
                .name(format!("timerwheel-worker-{index}"))
                .spawn(move || worker::run(worker_receiver, worker_metrics, panic_handler))
                .map_err(|_| Error::InvalidConfig("failed to start worker"))?;
            workers.push(handle);
        }

        Ok(Self {
            sender,
            workers: Mutex::new(workers),
            metrics,
            reject_policy: builder.reject_policy,
            closed: AtomicBool::new(false),
        })
    }

    /// Attempts to enqueue a closure without waiting for queue capacity.
    pub fn try_execute<F>(&self, task: F) -> Result<()>
    where
        F: FnOnce() + Send + 'static,
    {
        self.try_execute_box(Box::new(task))
    }

    pub(crate) fn try_execute_box(&self, task: BoxTask) -> Result<()> {
        self.try_execute_task(task).map_err(|rejected| {
            let (error, _) = rejected.into_parts();
            error
        })
    }

    fn try_execute_task(&self, task: BoxTask) -> std::result::Result<(), RejectedTask> {
        if self.closed.load(Ordering::Acquire) {
            return Err(RejectedTask::new(Error::Closed, task));
        }

        match self.sender.try_send(Message::Run(task)) {
            Ok(()) => {
                self.metrics.submitted();
                Ok(())
            }
            Err(mpsc::TrySendError::Full(task)) => {
                self.metrics.rejected();
                // Return the original boxed task so schedulers can retry
                // dispatch without losing ownership of the work.
                Err(RejectedTask::new(
                    Error::Saturated,
                    match task {
                        Message::Run(task) => task,
                        Message::Stop => {
                            unreachable!("task submission cannot create stop messages")
                        }
                    },
                ))
            }
            Err(mpsc::TrySendError::Disconnected(task)) => {
                self.metrics.rejected();
                Err(RejectedTask::new(
                    Error::Closed,
                    match task {
                        Message::Run(task) => task,
                        Message::Stop => {
                            unreachable!("task submission cannot create stop messages")
                        }
                    },
                ))
            }
        }
    }

    /// Enqueues a closure according to the configured reject policy.
    pub fn execute<F>(&self, task: F) -> Result<()>
    where
        F: FnOnce() + Send + 'static,
    {
        self.execute_box(Box::new(task))
    }

    pub(crate) fn execute_box(&self, task: BoxTask) -> Result<()> {
        if self.closed.load(Ordering::Acquire) {
            return Err(Error::Closed);
        }

        match self.reject_policy {
            RejectPolicy::Reject => self.try_execute_box(task),
            RejectPolicy::Block => match self.sender.send(Message::Run(task)) {
                Ok(()) => {
                    self.metrics.submitted();
                    Ok(())
                }
                Err(_) => Err(Error::Closed),
            },
        }
    }

    /// Stops accepting work and waits for worker threads to exit.
    pub fn shutdown(&self) -> Result<()> {
        if self.closed.swap(true, Ordering::AcqRel) {
            return Ok(());
        }

        // Stop messages are enqueued after accepted work, so workers drain
        // queued tasks before exiting.
        let worker_count = self.metrics.snapshot().workers;
        for _ in 0..worker_count {
            let _ = self.sender.send(Message::Stop);
        }

        if let Ok(mut workers) = self.workers.lock() {
            while let Some(worker) = workers.pop() {
                let _ = worker.join();
            }
        }

        Ok(())
    }

    /// Returns an immutable metrics snapshot.
    pub fn metrics(&self) -> ExecutorMetrics {
        self.metrics.snapshot()
    }
}

impl Executor for Pool {
    fn try_execute(&self, task: BoxTask) -> std::result::Result<(), RejectedTask> {
        self.try_execute_task(task)
    }

    fn shutdown(&self) -> Result<()> {
        Pool::shutdown(self)
    }

    fn metrics(&self) -> ExecutorMetrics {
        Pool::metrics(self)
    }
}

impl Drop for Pool {
    fn drop(&mut self) {
        let _ = self.shutdown();
    }
}

impl fmt::Debug for Pool {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Pool")
            .field("metrics", &self.metrics())
            .field("reject_policy", &self.reject_policy)
            .field("closed", &self.closed.load(Ordering::Acquire))
            .finish_non_exhaustive()
    }
}