mosaik 0.4.2

use {
	super::{
		super::{FunctionId, accept::CallReply},
		Stats,
	},
	crate::{
		discovery::{Discovery, PeerInfo},
		network::LocalNode,
		primitives::{Bytes, ShortFmtExt},
		tickets::TicketValidator,
	},
	dashmap::{DashMap, Entry as MapEntry},
	futures::future::BoxFuture,
	std::sync::Arc,
	tokio::sync::Semaphore,
	tokio_util::sync::CancellationToken,
};

/// A type-erased dispatch function that decodes a call payload, runs the
/// registered function, and encodes the reply.
pub(in crate::functions) type DispatchFn = Arc<
	dyn Fn(
			Bytes,
			crate::functions::CallContext,
		) -> BoxFuture<'static, Result<CallReply, DispatchError>>
		+ Send
		+ Sync,
>;

/// Handler-side dispatch faults. These are distinct from application-level
/// errors returned by the function itself, which travel back to the caller
/// as a successful [`CallReply::Err`] exchange.
#[derive(Debug, Clone, Copy)]
pub(in crate::functions) enum DispatchError {
	/// The call payload could not be decoded into the function's request
	/// type. This indicates a protocol violation by the caller.
	DecodeRequest,

	/// The function's reply (or application error) could not be encoded.
	/// This indicates a handler-side fault.
	EncodeReply,
}

/// A registered function handler entry.
pub(in crate::functions) struct Entry {
	/// The type-erased invocation path for the registered function.
	pub dispatch: DispatchFn,

	/// Predicate determining whether a caller peer is allowed to invoke
	/// this function.
	pub require: Box<dyn Fn(&PeerInfo) -> bool + Send + Sync>,

	/// Ticket validators that caller peers must satisfy.
	pub caller_auth: Vec<Arc<dyn TicketValidator>>,

	/// Bounds the number of concurrently executing calls.
	pub permits: Arc<Semaphore>,

	/// Cancelled when the owning [`Handler`](super::Handler) handle is
	/// dropped, aborting in-flight dispatches.
	pub cancel: CancellationToken,

	/// Invocation statistics for this function.
	pub stats: Arc<Stats>,
}

/// Responsible for keeping track of all locally served functions.
///
/// Each registered function is advertised in the local peer's discovery
/// entry so remote callers can find it, and un-advertised when its handler
/// is dropped.
pub(in crate::functions) struct Registry {
	/// The local node instance of the network associated with this
	/// functions subsystem.
	pub local: LocalNode,

	/// The discovery system used to announce registered functions.
	pub discovery: Discovery,

	/// Map of all registered function handlers by function id.
	pub active: DashMap<FunctionId, Arc<Entry>>,
}

impl Registry {
	/// Creates a new empty registry.
	pub(in crate::functions) fn new(
		local: LocalNode,
		discovery: Discovery,
	) -> Self {
		Self {
			local,
			discovery,
			active: DashMap::new(),
		}
	}

	/// Registers a new function handler entry for the given function id and
	/// advertises the function in the local discovery entry.
	///
	/// Returns an error if a handler for this function id is already
	/// registered.
	pub fn create(
		&self,
		function_id: FunctionId,
		entry: Entry,
	) -> Result<Arc<Entry>, ()> {
		match self.active.entry(function_id) {
			MapEntry::Vacant(slot) => {
				let entry = slot.insert(Arc::new(entry)).clone();
				let labels = [("network", self.local.network_id().short().to_string())];
				metrics::gauge!("mosaik.functions.handlers.active", &labels)
					.increment(1.0);

				// Update our local peer entry in discovery to include this function
				// id in the list of advertised functions so it can be discovered by
				// others.
				self
					.discovery
					.update_local_entry(move |me| me.add_functions(function_id));

				Ok(entry)
			}
			MapEntry::Occupied(_) => Err(()),
		}
	}

	/// Opens an existing handler entry for the given function id, if it
	/// exists.
	pub fn open(&self, function_id: FunctionId) -> Option<Arc<Entry>> {
		self.active.get(&function_id).map(|entry| entry.clone())
	}

	/// Deregisters the handler for the given function id and removes the
	/// function from the local discovery entry.
	pub fn remove(&self, function_id: FunctionId) {
		if self.active.remove(&function_id).is_some() {
			let labels = [("network", self.local.network_id().short().to_string())];
			metrics::gauge!("mosaik.functions.handlers.active", &labels)
				.decrement(1.0);

			self
				.discovery
				.update_local_entry(move |me| me.remove_functions(function_id));
		}
	}
}

impl core::fmt::Debug for Registry {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		write!(f, "Registry({} functions)", self.active.len())
	}
}