tightbeam-rs 0.9.0

//! Cluster framework for servlet orchestration
//!
//! Clusters are gateways that receive work requests from external clients
//! and route them to registered hives/hives based on servlet type.
//!
//! # Architecture
//!
//! 1. **Hives/Drones register** with the cluster, announcing available servlet types
//! 2. **Cluster maintains registry** of hives and their capabilities
//! 3. **Clients send** `ClusterWorkRequest` with `servlet_type` and `payload`
//! 4. **Cluster routes** to a hive that supports the requested servlet type
//! 5. **Cluster forwards** payload and returns response to client

pub mod builder;
pub mod error;
pub mod macros;
pub mod registry;
pub mod servlet_registry;

// Re-export submodule types
pub use builder::{ClusterConfBuilder, HeartbeatConfBuilder};
pub use error::ClusterError;
pub use registry::{HiveEntry, HiveRegistry, SharedId};
pub use servlet_registry::{PheromoneConf, ServletEntry, ServletRegistry};

use core::future::Future;
use core::marker::PhantomData;
use core::time::Duration;
use std::sync::Arc;

use crate::crypto::hash::{Digest, Sha3_256};
use crate::crypto::key::SigningKeyProvider;
use crate::policy::GatePolicy;
use crate::trace::TraceCollector;
use crate::transport::client::pool::PoolConfig;
use crate::transport::{Protocol, TightBeamAddress};

#[cfg(feature = "x509")]
use crate::crypto::x509::{policy::CertificateValidation, CertificateSpec};

use super::common::LeastLoaded;
use super::hive::LoadBalancer;

// =============================================================================
// Configuration
// =============================================================================

// Heartbeat default constants (single source of truth)
pub(crate) const DEFAULT_HEARTBEAT_INTERVAL_SECS: u64 = 5;
pub(crate) const DEFAULT_HEARTBEAT_TIMEOUT_SECS: u64 = 15;
pub(crate) const DEFAULT_MAX_CONCURRENT: usize = 10;
pub(crate) const DEFAULT_MAX_FAILURES: u32 = 3;

/// Configuration for cluster heartbeat behavior
///
/// Retry semantics are expressed through `interval` (cadence) and
/// `max_failures` (tolerance): a failed heartbeat is retried on the next
/// cycle rather than through a separate retry policy.
pub struct HeartbeatConf {
	/// Interval between heartbeat checks
	pub interval: Duration,
	/// Timeout before evicting unresponsive hives
	pub timeout: Duration,
	/// Maximum concurrent heartbeat requests
	pub max_concurrent: usize,
	/// Failed heartbeats before eviction
	pub max_failures: u32,
	/// Optional callback for heartbeat events (monitoring, testing)
	pub on_heartbeat: Option<HeartbeatCallback>,
}

impl Default for HeartbeatConf {
	fn default() -> Self {
		Self {
			interval: Duration::from_secs(DEFAULT_HEARTBEAT_INTERVAL_SECS),
			timeout: Duration::from_secs(DEFAULT_HEARTBEAT_TIMEOUT_SECS),
			max_concurrent: DEFAULT_MAX_CONCURRENT,
			max_failures: DEFAULT_MAX_FAILURES,
			on_heartbeat: None,
		}
	}
}

impl core::fmt::Debug for HeartbeatConf {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		f.debug_struct("HeartbeatConf")
			.field("interval", &self.interval)
			.field("timeout", &self.timeout)
			.field("max_concurrent", &self.max_concurrent)
			.field("max_failures", &self.max_failures)
			.field("on_heartbeat", &self.on_heartbeat.as_ref().map(|_| "Some(...)"))
			.finish()
	}
}

impl HeartbeatConf {
	/// Add a callback to be invoked on each heartbeat result
	pub fn with_callback(mut self, callback: HeartbeatCallback) -> Self {
		self.on_heartbeat = Some(callback);
		self
	}
}

// =============================================================================
// Heartbeat Callback
// =============================================================================

/// Event emitted for each heartbeat result
///
/// Provides information about the heartbeat outcome for monitoring,
/// metrics collection, or testing purposes.
#[derive(Debug, Clone)]
pub struct HeartbeatEvent {
	/// Address of the hive that was checked
	pub hive_addr: Arc<[u8]>,
	/// Whether the heartbeat succeeded
	pub success: bool,
	/// Utilization reported by the hive (if successful)
	pub utilization: Option<crate::utils::BasisPoints>,
}

/// Callback type for heartbeat events
///
/// Called after each heartbeat result is processed. The callback must be
/// thread-safe (`Send + Sync`) as it may be invoked from multiple concurrent
/// heartbeat tasks.
pub type HeartbeatCallback = Arc<dyn Fn(HeartbeatEvent) + Send + Sync>;

// ============================================================================
// TLS Configuration
// ============================================================================

/// TLS configuration for cluster -> hive connections
///
/// Contains certificate, key, and validators for encrypted transport.
/// Used by the connection pool for mutual TLS with hives.
#[cfg(feature = "x509")]
pub struct ClusterTlsConfig {
	/// Cluster's certificate (used for both client and hive connections)
	pub certificate: CertificateSpec,
	/// Private key provider for signing operations (supports HSM/KMS)
	pub key: Arc<dyn SigningKeyProvider>,
	/// Server certificate validators for hive connections (cluster->hive)
	pub validators: Vec<Arc<dyn CertificateValidation>>,
	/// Client certificate validators for mutual auth (client->cluster)
	pub client_validators: Vec<Arc<dyn CertificateValidation>>,
	/// Trust store for validating hive/servlet server certificates (outbound connections)
	pub hive_trust: Option<Arc<dyn crate::crypto::x509::store::CertificateTrust>>,
}

#[cfg(feature = "x509")]
impl Clone for ClusterTlsConfig {
	fn clone(&self) -> Self {
		Self {
			certificate: self.certificate.clone(),
			key: Arc::clone(&self.key),
			validators: self.validators.iter().map(Arc::clone).collect(),
			client_validators: self.client_validators.iter().map(Arc::clone).collect(),
			hive_trust: self.hive_trust.as_ref().map(Arc::clone),
		}
	}
}

#[cfg(feature = "x509")]
impl core::fmt::Debug for ClusterTlsConfig {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		f.debug_struct("ClusterTlsConfig")
			.field("certificate", &self.certificate)
			.field("key", &"<KeyProvider>")
			.field("validators", &format!("[{} validators]", self.validators.len()))
			.field("client_validators", &format!("[{} validators]", self.client_validators.len()))
			.field("hive_trust", &self.hive_trust.as_ref().map(|_| "Some(<TrustStore>)"))
			.finish()
	}
}

/// Configuration for clusters
///
/// Contains settings for load balancing, health checks, gateway policies,
/// and cryptographic signing for cluster -> hive communication.
///
/// # Type Parameters
/// - `L`: Load balancing strategy (default: `LeastLoaded`)
/// - `D`: Digest algorithm for frame integrity and signing (default: `Sha3_256`)
pub struct ClusterConf<L: LoadBalancer = LeastLoaded, D: Digest = Sha3_256> {
	/// Load balancing strategy for distributing work across hives
	pub load_balancer: L,
	/// Heartbeat configuration
	pub heartbeat: HeartbeatConf,
	/// Pheromone configuration for bio-inspired routing
	pub pheromone: PheromoneConf,
	/// Gate policies for the gateway (rate limiting, auth, etc.)
	pub policies: Vec<Arc<dyn GatePolicy + Send + Sync>>,
	/// Connection pool configuration for hive connections
	pub pool_config: PoolConfig,
	/// Freshness window in milliseconds for signed hive control frames
	/// (registration, address updates); stale or replayed frames inside
	/// the window are rejected (CWE-294)
	pub control_freshness_window_ms: u64,
	/// TLS configuration for cluster -> hive connections
	#[cfg(feature = "x509")]
	pub tls: ClusterTlsConfig,
	/// Phantom data for digest type
	pub(crate) _digest: PhantomData<D>,
}

#[cfg(feature = "x509")]
impl ClusterConf {
	/// Create a new cluster configuration with TLS config
	pub fn new(tls: ClusterTlsConfig) -> Self {
		Self::builder(tls).build()
	}
}

#[cfg(feature = "x509")]
impl<L: LoadBalancer, D: Digest> core::fmt::Debug for ClusterConf<L, D> {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		f.debug_struct("ClusterConfig")
			.field("heartbeat", &self.heartbeat)
			.field("pheromone", &self.pheromone)
			.field("policies", &format!("[{} policies]", self.policies.len()))
			.field("pool_config", &self.pool_config)
			.field("control_freshness_window_ms", &self.control_freshness_window_ms)
			.field("tls", &self.tls)
			.finish()
	}
}

// =============================================================================
// Work Request/Response Messages
// =============================================================================

pub use crate::colony::common::{ClusterRequest, ClusterWorkRequest, ClusterWorkResponse};

// =============================================================================
// Cluster Trait
// =============================================================================

/// Trait for cluster implementations
///
/// Clusters are gateways that route work requests to registered hives
/// based on servlet type. Hives register dynamically, and the cluster
/// learns available servlet types from their registrations.
pub trait Cluster: Sized + Send + Sync {
	/// The protocol type this cluster uses
	type Protocol: Protocol;

	/// Address type for this cluster
	type Address: TightBeamAddress;

	/// Start the cluster gateway
	fn start(
		trace: Arc<TraceCollector>,
		config: ClusterConf,
	) -> impl Future<Output = Result<Self, crate::TightBeamError>> + Send;

	/// Get the gateway address
	fn addr(&self) -> Self::Address;

	/// Get available servlet types (from registered hives)
	fn available_servlets(&self) -> Vec<Vec<u8>>;

	/// Get the number of registered hives
	fn hive_count(&self) -> usize;

	/// Get the trace collector
	fn trace(&self) -> Arc<TraceCollector>;

	/// Stop the cluster
	fn stop(self);

	/// Wait for the cluster to finish
	fn join(self) -> impl Future<Output = Result<(), crate::colony::servlet::servlet_runtime::rt::JoinError>> + Send;

	// =========================================================================
	// Heartbeat Methods
	// =========================================================================

	/// Access the hive registry
	fn registry(&self) -> &Arc<HiveRegistry>;

	/// Access heartbeat configuration
	fn heartbeat_config(&self) -> &HeartbeatConf;

	/// Send a single heartbeat to a hive
	///
	/// Builds a signed heartbeat frame and sends it via the connection pool.
	/// Returns the heartbeat result from the hive.
	///
	/// The heartbeat loop itself is generated by the `cluster!` macro
	/// (tokio `JoinSet` with bounded concurrency); it is not part of the
	/// trait surface.
	fn send_heartbeat(
		&self,
		addr: Self::Address,
	) -> impl Future<Output = Result<super::common::HeartbeatResult, ClusterError>> + Send;
}

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

#[cfg(test)]
mod tests {
	use super::*;
	use crate::colony::common::RegisterHiveRequest;
	use crate::colony::hive::ServletInfo;
	use crate::crypto::key::Secp256k1KeyProvider;
	use crate::crypto::sign::ecdsa::Secp256k1SigningKey;
	use crate::policy::TransitStatus;
	use crate::testing::create_test_signing_key;
	use crate::utils::BasisPoints;

	// =========================================================================
	// Test Helpers
	// =========================================================================

	fn test_tls_config() -> ClusterTlsConfig {
		let key: Secp256k1SigningKey = create_test_signing_key();
		ClusterTlsConfig {
			certificate: CertificateSpec::Der(&[]),
			key: Arc::new(Secp256k1KeyProvider::from(key)),
			validators: Vec::new(),
			client_validators: Vec::new(),
			hive_trust: None,
		}
	}

	fn test_registry() -> HiveRegistry {
		HiveRegistry::new(Duration::from_secs(15))
	}

	fn request(addr: &[u8], servlets: &[&[u8]]) -> RegisterHiveRequest {
		RegisterHiveRequest {
			issued_at_ms: 0,
			hive_addr: addr.to_vec(),
			metadata: None,
			servlet_addresses: servlets
				.iter()
				.map(|s| ServletInfo { servlet_id: s.to_vec(), address: addr.to_vec() })
				.collect(),
		}
	}

	fn request_with_meta(addr: &[u8], servlets: &[&[u8]], meta: &[u8]) -> RegisterHiveRequest {
		RegisterHiveRequest {
			issued_at_ms: 0,
			hive_addr: addr.to_vec(),
			metadata: Some(meta.to_vec()),
			servlet_addresses: servlets
				.iter()
				.map(|s| ServletInfo { servlet_id: s.to_vec(), address: addr.to_vec() })
				.collect(),
		}
	}

	// =========================================================================
	// ClusterConf Tests
	// =========================================================================

	#[test]
	fn cluster_conf_defaults() {
		let config = ClusterConf::new(test_tls_config());
		assert_eq!(config.heartbeat.interval, Duration::from_secs(5));
		assert_eq!(config.heartbeat.timeout, Duration::from_secs(15));
		assert!(config.policies.is_empty());
	}

	// =========================================================================
	// ClusterWorkResponse Tests
	// =========================================================================

	#[test]
	fn work_response_ok() {
		let response = ClusterWorkResponse::ok(b"test".to_vec());
		assert_eq!(response.status, TransitStatus::Accepted);
		assert_eq!(response.payload, Some(b"test".to_vec()));
	}

	#[test]
	fn work_response_err() {
		let response = ClusterWorkResponse::err(TransitStatus::Forbidden);
		assert_eq!(response.status, TransitStatus::Forbidden);
		assert!(response.payload.is_none());
	}

	// =========================================================================
	// HiveRegistry Tests
	// =========================================================================

	#[test]
	fn registry_register_and_lookup() -> Result<(), ClusterError> {
		let registry = test_registry();
		registry.register(request(b"127.0.0.1:8080", &[b"ping", b"calc"]))?;

		// Registered types found
		assert_eq!(registry.hives_for_type(b"ping")?.len(), 1);
		assert_eq!(registry.hives_for_type(b"calc")?.len(), 1);
		assert_eq!(registry.hives_for_type(b"ping")?[0].address.as_ref(), b"127.0.0.1:8080");

		// Unknown type not found
		assert!(registry.hives_for_type(b"unknown")?.is_empty());

		Ok(())
	}

	#[test]
	fn registry_unregister() -> Result<(), ClusterError> {
		let registry = test_registry();
		registry.register(request(b"127.0.0.1:8080", &[b"ping"]))?;

		assert_eq!(registry.len()?, 1);
		assert!(registry.unregister(b"127.0.0.1:8080")?.is_some());
		assert_eq!(registry.len()?, 0);
		assert!(registry.hives_for_type(b"ping")?.is_empty());

		Ok(())
	}

	#[test]
	fn registry_update_utilization() -> Result<(), ClusterError> {
		let registry = test_registry();
		registry.register(request(b"127.0.0.1:8080", &[b"ping"]))?;

		assert!(registry.update_utilization(b"127.0.0.1:8080", BasisPoints::new(5000))?);
		assert_eq!(registry.hives_for_type(b"ping")?[0].utilization.get(), 5000);

		Ok(())
	}

	#[test]
	fn registry_available_servlets_deduplicated() -> Result<(), ClusterError> {
		let registry = test_registry();
		registry.register(request(b"hive1", &[b"ping", b"calc"]))?;
		registry.register(request(b"hive2", &[b"ping", b"worker"]))?;

		// ping, calc, worker - ping deduplicated
		assert_eq!(registry.to_available_servlets()?.len(), 3);

		Ok(())
	}

	#[test]
	fn registry_multiple_hives_same_type() -> Result<(), ClusterError> {
		let registry = test_registry();
		registry.register(request(b"hive1", &[b"ping"]))?;
		registry.register(request(b"hive2", &[b"ping"]))?;

		assert_eq!(registry.hives_for_type(b"ping")?.len(), 2);

		Ok(())
	}

	#[test]
	fn registry_all_hives() -> Result<(), ClusterError> {
		let registry = test_registry();
		registry.register(request(b"hive1", &[b"ping"]))?;
		registry.register(request_with_meta(b"hive2", &[b"calc"], b"metadata"))?;

		let all = registry.all_hives()?;
		assert_eq!(all.len(), 2);

		let addrs: Vec<_> = all.iter().map(|e| e.address.as_ref()).collect();
		assert!(addrs.contains(&b"hive1".as_slice()));
		assert!(addrs.contains(&b"hive2".as_slice()));

		Ok(())
	}
}