use std::{
collections::HashMap,
error::Error,
task::{Context, Poll},
};
use futures::future::{self, BoxFuture};
use snafu::Snafu;
use crate::server::UnresolvedRequest;
#[derive(Debug, Clone, Snafu)]
#[snafu(visibility(pub))]
pub enum ServersRouterDispatchError {
#[snafu(display("service not found for server name: {server_name}"))]
MissingService { server_name: String },
}
#[derive(Debug, Clone)]
pub struct ServersRouter<S> {
router: std::sync::Arc<HashMap<String, S>>,
}
impl<S> Default for ServersRouter<S> {
fn default() -> Self {
Self {
router: Default::default(),
}
}
}
impl<S> ServersRouter<S> {
pub fn new() -> Self {
Self::default()
}
pub fn contains(&self, domain: &str) -> bool {
self.router.contains_key(domain)
}
}
impl<S: Clone> ServersRouter<S> {
fn router_mut(&mut self) -> &mut HashMap<String, S> {
std::sync::Arc::make_mut(&mut self.router)
}
pub fn insert(&mut self, domain: impl Into<String>, service: S) -> Option<S> {
self.router_mut().insert(domain.into(), service)
}
pub fn serve(&mut self, domain: impl Into<String>, service: S) -> &mut Self {
_ = self.insert(domain, service);
self
}
}
impl<S> tower_service::Service<UnresolvedRequest> for ServersRouter<S>
where
S: tower_service::Service<UnresolvedRequest, Response = ()> + Clone + Send + Sync + 'static,
S::Future: Send,
S::Error: Into<Box<dyn Error + Send + Sync>>,
{
type Response = ();
type Error = Box<dyn Error + Send + Sync>;
type Future = BoxFuture<'static, Result<(), Self::Error>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, req: UnresolvedRequest) -> Self::Future {
let server_name = req.local_agent.name().to_string();
let Some(mut service) = self.router.get(server_name.as_str()).cloned() else {
return Box::pin(async move {
Err(ServersRouterDispatchError::MissingService { server_name }.into())
});
};
Box::pin(async move {
future::poll_fn(|cx| service.poll_ready(cx))
.await
.map_err(Into::into)?;
service.call(req).await.map_err(Into::into)
})
}
}