tightbeam-rs 0.9.0

//! Worker framework for message processing
//!
//! Workers are the fundamental processing units in the colony architecture.
//! They receive messages, apply policies, and produce responses.
//!
//! # Example
//!
//! ```ignore
//! worker! {
//!     name: MyWorker<RequestMessage, ResponseMessage>,
//!     config: { threshold: u32 },
//!     policies: { with_receptor_gate: [MyGate] },
//!     handle: |message, trace, config| async move {
//!         ResponseMessage { value: message.value + config.threshold }
//!     }
//! }
//! ```

#[cfg(not(feature = "std"))]
extern crate alloc;
#[cfg(not(feature = "std"))]
use alloc::boxed::Box;
#[cfg(not(feature = "std"))]
use alloc::sync::Arc;

#[cfg(feature = "std")]
use std::sync::Arc;

use core::future::Future;
use core::pin::Pin;

#[cfg(feature = "derive")]
use crate::Errorizable;

use crate::policy::{ReceptorPolicy, TransitStatus};
use crate::trace::TraceCollector;
use crate::Message;

pub mod macros;

/// Re-export unified runtime primitives with worker-specific type aliases
pub mod worker_runtime {
	pub mod rt {
		pub use crate::runtime::rt::*;

		/// Queue sender type alias (for backwards compatibility)
		pub type QueueSender<T> = crate::runtime::rt::Sender<T>;

		/// Queue receiver type alias (for backwards compatibility)
		pub type QueueReceiver<T> = crate::runtime::rt::Receiver<T>;
	}
}

pub struct WorkerRequest<I: Send, O> {
	pub message: Arc<I>,
	pub respond_to: worker_runtime::rt::OneshotSender<Result<O, TransitStatus>>,
	pub trace: Arc<TraceCollector>,
}

#[cfg_attr(feature = "derive", derive(Errorizable))]
#[derive(Debug)]
pub enum WorkerRelayError {
	#[cfg_attr(feature = "derive", error("Worker queue closed"))]
	QueueClosed,
	#[cfg_attr(feature = "derive", error("Worker response channel dropped"))]
	ResponseDropped,
	#[cfg_attr(feature = "derive", error("Message rejected with status {:?}"))]
	#[cfg_attr(feature = "derive", from)]
	Rejected(TransitStatus),
}

#[cfg(not(feature = "derive"))]
impl core::fmt::Display for WorkerRelayError {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		match self {
			Self::QueueClosed => f.write_str("worker queue closed"),
			Self::ResponseDropped => f.write_str("worker response channel dropped"),
			Self::Rejected(status) => write!(f, "message rejected with status {:?}", status),
		}
	}
}

#[cfg(not(feature = "derive"))]
impl std::error::Error for WorkerRelayError {}

pub type WorkerRelayFuture<O> = Pin<Box<dyn Future<Output = Result<O, WorkerRelayError>> + Send + 'static>>;
pub type WorkerStartFuture<W> = Pin<Box<dyn Future<Output = Result<W, crate::error::TightBeamError>> + Send>>;
pub type WorkerKillFuture = Pin<Box<dyn Future<Output = Result<(), crate::error::TightBeamError>> + Send + 'static>>;

/// Relay a message into a started worker's queue and await its response
///
/// Single implementation backing every `worker!`-generated `Worker::relay`.
pub fn relay_to_worker<I, O>(
	sender: Option<worker_runtime::rt::QueueSender<WorkerRequest<I, O>>>,
	trace: Arc<TraceCollector>,
	message: Arc<I>,
) -> WorkerRelayFuture<O>
where
	I: Send + Sync + 'static,
	O: Send + 'static,
{
	Box::pin(async move {
		let sender = sender.ok_or(WorkerRelayError::QueueClosed)?;
		let (tx, rx) = worker_runtime::rt::oneshot();

		let request = WorkerRequest { message, respond_to: tx, trace };
		worker_runtime::rt::send(&sender, request)
			.await
			.map_err(|_| WorkerRelayError::QueueClosed)?;

		match worker_runtime::rt::wait_response(rx).await {
			Ok(Ok(output)) => Ok(output),
			Ok(Err(status)) => Err(WorkerRelayError::Rejected(status)),
			Err(()) => Err(WorkerRelayError::ResponseDropped),
		}
	})
}

/// Gracefully stop a started worker: close its queue, then await loop exit
///
/// Single implementation backing every `worker!`-generated `Worker::kill`.
/// Dropping the sender ends the run loop's `recv` stream, so the join
/// observes a clean exit instead of aborting mid-message.
pub fn kill_worker<I, O>(
	sender: Option<worker_runtime::rt::QueueSender<WorkerRequest<I, O>>>,
	join: Option<worker_runtime::rt::JoinHandle>,
) -> WorkerKillFuture
where
	I: Send + Sync + 'static,
	O: Send + 'static,
{
	Box::pin(async move {
		drop(sender);

		if let Some(handle) = join {
			worker_runtime::rt::join(handle)
				.await
				.map_err(|_| crate::error::TightBeamError::JoinError)?;
		}

		Ok(())
	})
}

/// A message-processing unit in the colony architecture
///
/// A worker owns a bounded request queue and a spawned run loop that
/// evaluates receptor policies, invokes the handler, and responds through a
/// oneshot channel per request. Implementations are generated by the
/// [`worker!`](crate::worker) macro; servlets look workers up by their
/// [`WorkerMetadata::name`].
///
/// # Lifecycle
///
/// [`new`](Worker::new) -> [`start`](Worker::start) ->
/// [`relay`](Worker::relay) (any number of times) -> [`kill`](Worker::kill).
/// Dropping a started worker without killing it aborts the run loop instead
/// of draining it.
pub trait Worker: Send + Sync + Sized {
	/// Message type consumed by the handler
	type Input: Send + Sync + 'static;
	/// Response type produced by the handler
	type Output: Send + 'static;
	/// Construction-time configuration, exposed to the handler by
	/// reference; `()` for workers declared without a `config` block
	type Config: Send + Sync + 'static;

	/// Construct an unstarted worker: no queue, no run loop
	///
	/// [`relay`](Worker::relay) before [`start`](Worker::start) fails with
	/// [`WorkerRelayError::QueueClosed`].
	fn new(config: Self::Config) -> Self;

	/// Open the request queue and spawn the run loop
	///
	/// Consumes the unstarted worker and resolves to the started one.
	/// Starting an already-started worker is a no-op that returns it
	/// unchanged. The `trace` collector is attached to every relayed
	/// request.
	fn start(self, trace: Arc<TraceCollector>) -> WorkerStartFuture<Self>;

	/// Gracefully stop the worker
	///
	/// Closes the queue so the run loop drains in-flight requests and
	/// exits, then awaits the join. Fails with
	/// [`TightBeamError::JoinError`](crate::error::TightBeamError::JoinError)
	/// if the run loop task cannot be joined.
	fn kill(self) -> WorkerKillFuture;

	/// Enqueue a message and await the handler's response
	///
	/// Fails with [`WorkerRelayError::QueueClosed`] if the worker was never
	/// started or already killed, [`WorkerRelayError::Rejected`] if a
	/// receptor gate refuses the message, and
	/// [`WorkerRelayError::ResponseDropped`] if the run loop dies before
	/// responding.
	fn relay(&self, message: Arc<Self::Input>) -> WorkerRelayFuture<Self::Output>;

	/// Bound of the request queue
	///
	/// `0` before [`start`](Worker::start); afterwards the `queue:` value
	/// from the `worker!` declaration (default 64).
	fn queue_capacity(&self) -> usize;
}

/// Provides static metadata about a worker type
pub trait WorkerMetadata {
	/// Returns the registration name for this worker
	fn name() -> &'static str;
}

pub struct WorkerPolicies<I: Send> {
	#[allow(dead_code)]
	pub(crate) receptor_gates: Vec<Arc<dyn ReceptorPolicy<I> + Send + Sync>>,
}

impl<I: Send> WorkerPolicies<I> {
	pub fn receptor_gates(&self) -> &[Arc<dyn ReceptorPolicy<I> + Send + Sync>] {
		&self.receptor_gates
	}
}

impl<I: Send> Default for WorkerPolicies<I> {
	fn default() -> Self {
		Self { receptor_gates: Vec::new() }
	}
}

pub struct WorkerPolicyBuilder<I: Send> {
	receptor_gates: Vec<Arc<dyn ReceptorPolicy<I> + Send + Sync>>,
}

impl<I: Send> Default for WorkerPolicyBuilder<I> {
	fn default() -> Self {
		Self { receptor_gates: Vec::new() }
	}
}

impl<I: Message + Send> WorkerPolicyBuilder<I> {
	pub fn build(self) -> WorkerPolicies<I> {
		WorkerPolicies { receptor_gates: self.receptor_gates }
	}

	pub fn with_receptor_gate<R, const N: usize>(mut self, gates: [R; N]) -> Self
	where
		R: ReceptorPolicy<I> + Send + Sync + 'static,
	{
		self.receptor_gates.extend(
			gates
				.into_iter()
				.map(|gate| Arc::new(gate) as Arc<dyn ReceptorPolicy<I> + Send + Sync>),
		);
		self
	}
}

#[cfg(test)]
mod tests {
	use std::sync::Arc;

	use super::WorkerRelayError;
	use crate::der::Sequence;
	use crate::policy::{ReceptorPolicy, TransitStatus};
	use crate::worker;
	use crate::Beamable;

	#[derive(Beamable, Clone, Debug, PartialEq, Sequence)]
	pub struct RequestMessage {
		content: String,
		lucky_number: u32,
	}

	#[derive(Sequence, Beamable, Clone, Debug, PartialEq)]
	pub struct PongMessage {
		result: String,
	}

	#[derive(Default)]
	struct PingGate;

	impl ReceptorPolicy<RequestMessage> for PingGate {
		fn evaluate(&self, maybe_ping: &RequestMessage) -> TransitStatus {
			if maybe_ping.content == "PING" {
				TransitStatus::Accepted
			} else {
				TransitStatus::Forbidden
			}
		}
	}

	worker! {
		name: LuckyNumberDeterminer<RequestMessage, bool>,
		config: {
			lotto_number: u32,
		},
		handle: |message, _trace, config| async move {
			message.lucky_number == config.lotto_number
		}
	}

	worker! {
		name: PingPongWorker<RequestMessage, PongMessage>,
		policies: {
			with_receptor_gate: [PingGate]
		},
		handle: |_message, _trace| async move {
			PongMessage {
				result: "PONG".to_string(),
			}
		}
	}

	#[cfg(feature = "std")]
	crate::test_worker! {
		name: lucky_number_worker_checks_winner,
		setup: || {
			LuckyNumberDeterminer::new(LuckyNumberDeterminerConf { lotto_number: 42 })
		},
		assertions: |worker| async move {
			assert_eq!(worker.queue_capacity(), 64);

			let winner = worker.relay(Arc::new(RequestMessage {
				content: "PING".to_string(),
				lucky_number: 42,
			})).await?;
			assert!(winner);

			let loser = worker.relay(Arc::new(RequestMessage {
				content: "PING".to_string(),
				lucky_number: 7,
			})).await?;
			assert!(!loser);

			Ok(())
		}
	}

	#[cfg(feature = "std")]
	crate::test_worker! {
		name: test_ping_pong_worker,
		setup: || {
			PingPongWorker::new(())
		},
		assertions: |worker| async move {
			// Test accepted message
			let ping_msg = RequestMessage {
				content: "PING".to_string(),
				lucky_number: 42,
			};
			let response = worker.relay(Arc::new(ping_msg)).await?;
			assert_eq!(response, PongMessage { result: "PONG".to_string() });

			// Test rejected message
			let pong_msg = RequestMessage {
				content: "PONG".to_string(),
				lucky_number: 42,
			};

			let result = worker.relay(Arc::new(pong_msg)).await;
			assert!(matches!(result, Err(WorkerRelayError::Rejected(_))));

			Ok(())
		}
	}
}