serverless_fn/

server.rs

//! Server module for running serverless functions as HTTP services.

use std::future::Future;
use std::sync::Arc;

use futures::future::BoxFuture;

use crate::error::ServerlessError;

/// Function handler type.
pub type FunctionHandler =
    Arc<dyn Fn(Vec<u8>) -> BoxFuture<'static, Result<Vec<u8>, ServerlessError>> + Send + Sync>;

/// Unified trait for registering serverless functions.
///
/// Implementations should register themselves with the appropriate server type.
pub trait FunctionRegistry: Send + Sync + 'static {
    /// Returns the function name.
    fn function_name(&self) -> &'static str;

    /// Returns the function path (for HTTP routing).
    fn function_path(&self) -> &'static str {
        self.function_name()
    }

    /// Registers the function handler with the server.
    fn register(&self, server: &mut FunctionServer);
}

// Collect all FunctionRegistry implementations
inventory::collect!(&'static dyn FunctionRegistry);

/// Server configuration.
#[derive(Debug, Clone)]
pub struct ServerConfig {
    /// Host address to bind to.
    pub host: String,
    /// Port to listen on.
    pub port: u16,
}

impl Default for ServerConfig {
    fn default() -> Self {
        Self {
            host: "127.0.0.1".to_string(),
            port: 3000,
        }
    }
}

impl ServerConfig {
    /// Creates a new server configuration with default values.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the host address.
    #[must_use]
    pub fn host(mut self, host: &str) -> Self {
        self.host = host.to_string();
        self
    }

    /// Sets the port.
    #[must_use]
    pub fn port(mut self, port: u16) -> Self {
        self.port = port;
        self
    }
}

/// Unified server for hosting serverless functions.
///
/// # Examples
///
/// ```rust,no_run
/// use serverless_fn::server::{FunctionServer, ServerConfig};
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
///     // Method 1: Using builder pattern
///     FunctionServer::new()
///         .host("0.0.0.0")
///         .port(3000)
///         .start()
///         .await?;
///
///     // Method 2: Using custom config
///     let config = ServerConfig::new().host("0.0.0.0").port(3000);
///     FunctionServer::with_config(config)
///         .start()
///         .await?;
///
///     Ok(())
/// }
/// ```
pub struct FunctionServer {
    config: ServerConfig,
    router: Option<axum::Router>,
}

impl Default for FunctionServer {
    fn default() -> Self {
        Self::new()
    }
}

impl FunctionServer {
    /// Creates a new server with default configuration.
    #[must_use]
    pub fn new() -> Self {
        Self {
            config: ServerConfig::default(),
            router: Some(axum::Router::new()),
        }
    }

    /// Creates a new server with custom configuration.
    #[must_use]
    pub fn with_config(config: ServerConfig) -> Self {
        Self {
            config,
            router: Some(axum::Router::new()),
        }
    }

    /// Sets the host address to bind to.
    #[must_use]
    pub fn host(mut self, host: &str) -> Self {
        self.config.host = host.to_string();
        self
    }

    /// Sets the port to listen on.
    #[must_use]
    pub fn port(mut self, port: u16) -> Self {
        self.config.port = port;
        self
    }

    /// Returns a reference to the server configuration.
    #[must_use]
    pub fn config(&self) -> &ServerConfig {
        &self.config
    }

    /// Registers a function for HTTP transport.
    pub fn register_http_route<F, T>(&mut self, path: &str, handler: F)
    where
        F: axum::handler::Handler<T, ()> + Send + Sync + 'static,
        T: 'static,
    {
        if let Some(ref mut router) = self.router {
            let old_router = std::mem::replace(router, axum::Router::new());
            *router = old_router.route(path, axum::routing::post(handler));
        }
    }

    /// Starts the server and listens for incoming requests.
    ///
    /// Automatically discovers and registers all serverless functions
    /// using the inventory crate.
    ///
    /// # Errors
    ///
    /// Returns an error if the server fails to start.
    pub async fn start(mut self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        // Register all functions from inventory
        for registry in inventory::iter::<&'static dyn FunctionRegistry> {
            registry.register(&mut self);
        }

        self.serve().await
    }

    /// Starts the server with graceful shutdown support.
    ///
    /// # Errors
    ///
    /// Returns an error if the server fails to start.
    pub async fn start_with_graceful_shutdown<F>(
        mut self,
        shutdown_signal: F,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>
    where
        F: Future<Output = ()> + Send + 'static + Clone,
    {
        // Register all functions from inventory
        for registry in inventory::iter::<&'static dyn FunctionRegistry> {
            registry.register(&mut self);
        }

        self.serve_with_shutdown(shutdown_signal).await
    }

    /// Internal method to serve requests based on enabled features.
    async fn serve(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        if let Some(ref router) = self.router {
            let listener =
                tokio::net::TcpListener::bind((self.config.host.as_str(), self.config.port))
                    .await?;
            return axum::serve(listener, router.clone())
                .await
                .map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) });
        }

        Ok(())
    }

    /// Internal method to serve requests with graceful shutdown.
    async fn serve_with_shutdown<F>(
        &self,
        shutdown_signal: F,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>
    where
        F: Future<Output = ()> + Send + 'static,
    {
        if let Some(ref router) = self.router {
            let listener =
                tokio::net::TcpListener::bind((self.config.host.as_str(), self.config.port))
                    .await?;
            return axum::serve(listener, router.clone())
                .with_graceful_shutdown(shutdown_signal)
                .await
                .map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) });
        }

        Ok(())
    }
}

// Re-export for backward compatibility
pub use FunctionServer as Server;