rogue_runtime/

runtime.rs

use std::any::Any;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use anyhow::{Context, Result, anyhow};
use futures::future::BoxFuture;
#[cfg(any(feature = "client", feature = "server"))]
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use serde::{Deserialize, Serialize};
use serde_binary::binary_stream::Endian;
#[cfg(feature = "server")]
use tokio::net::{TcpListener, TcpStream};
use tokio::spawn;
use tokio::sync::{Mutex, RwLock, mpsc};
use tokio::task::JoinHandle;
#[cfg(feature = "client")]
use tokio_tungstenite::connect_async;
#[cfg(any(feature = "client", feature = "server"))]
use tokio_tungstenite::tungstenite;
use tracing::{error, info, trace};
use ulid::Ulid;

#[cfg(feature = "client")]
use crate::message::HandshakeMessage;
use crate::message::{Message, MessageBody, RpcMessage};
use crate::{
    CONTEXT, Identity, InstanceId, InventoryItem, MessageHandler, MessageId, ObjectRef,
    RuntimeContext, RuntimeId, TypeId, TypeInfo,
};

#[cfg(feature = "server")]
#[allow(unused_imports)]
use crate::api;

/// Type alias for a pinned, boxed future that resolves to a [`Message`].
type MessageFuture = Pin<Box<dyn Future<Output = Message> + Send>>;

/// Type alias for an asynchronous executor that takes a [`Message`] and returns a `MessageFuture`.
type Executor = Box<dyn Fn(Message) -> MessageFuture + Send + Sync>;

/// Internal envelope pairing an incoming [`Message`] with the channel used to send the response back
/// to the caller.
struct MessageData {
    message: Message,
    return_tx: mpsc::Sender<Message>,
}

/// Information about runtime
#[derive(Clone, Deserialize, Serialize)]
pub struct RuntimeInfo {
    pub id: RuntimeId,
}

#[crate::async_trait]
pub trait RuntimeTrait {
    async fn create(id: RuntimeId) -> Result<Self>
    where
        Self: Sized;
    async fn execute<F>(&self, target_id: RuntimeId, f: F) -> Result<()>
    where
        F: Future<Output = Result<()>> + Send + 'static;
    async fn execute_local<F>(&self, f: F) -> Result<()>
    where
        F: Future<Output = Result<()>> + Send + 'static;
    async fn register_handler<H: MessageHandler>(&self) -> Result<()>;
    async fn register_instance<T>(&self, instance: T) -> ObjectRef<T>
    where
        T: Identity + Any + Send + Sync + 'static;
    async fn get_instance<T>(&self, id: InstanceId) -> Result<Arc<RwLock<T>>, String>
    where
        T: Identity + Any + Send + Sync + 'static;
    async fn take_instance<T>(&self, id: InstanceId) -> Option<T>
    where
        T: Identity + Any + Send + Sync + 'static;
    async fn call<T: MessageHandler>(
        &self,
        target_id: RuntimeId,
        args: T::Input,
    ) -> Result<T::Output>;
    #[cfg(any(feature = "client", feature = "server"))]
    async fn connected_runtimes(&self) -> Vec<RuntimeInfo>;
    async fn inventory(&self) -> Vec<InventoryItem>;
    #[cfg(feature = "client")]
    async fn connect(&self, addr: String) -> Result<JoinHandle<()>>;
    #[cfg(feature = "server")]
    async fn start_server(&self, addr: String) -> JoinHandle<Result<()>>;
}

#[crate::async_trait]
pub(crate) trait RuntimeInternalTrait {
    async fn register_handlers(&self) -> Result<()>;
    async fn runtime_worker(&self);
    #[cfg(any(feature = "client", feature = "server"))]
    async fn message_encoder<S: Sink<tungstenite::Message> + Unpin + Send>(
        tx: S,
        rx: mpsc::Receiver<Message>,
    ) where
        <S as Sink<tungstenite::Message>>::Error: std::fmt::Display;
    #[cfg(any(feature = "client", feature = "server"))]
    async fn message_decoder<
        S: Stream<Item = std::result::Result<tungstenite::Message, tungstenite::Error>> + Unpin + Send,
    >(
        &self,
        tx: mpsc::Sender<Message>,
        rx: S,
    );
    #[cfg(feature = "server")]
    async fn accept_connection(&self, stream: TcpStream) -> Result<()>;
}

/// Registration entry for automatic handler registration via inventory
pub struct HandlerRegistration {
    /// Given a runtime instance, returns a future that registers one handler
    pub register: fn(Runtime) -> BoxFuture<'static, Result<()>>,
}

// Collect all registrations submitted by the `rpc` macro
inventory::collect!(HandlerRegistration);

pub type Runtime = Arc<RuntimeImpl>;

/// The core RPC runtime.
///
/// Responsible for sending and receiving messages, dispatching
/// requests to registered handlers, and managing message executors.
pub struct RuntimeImpl {
    id: RuntimeId,
    executors: Arc<RwLock<HashMap<TypeId, Arc<Executor>>>>,
    /// Stores return channels for in-flight RPC calls
    message_handlers: Arc<Mutex<HashMap<Ulid, mpsc::Sender<Message>>>>,
    inventory: Arc<Mutex<HashMap<TypeId, InventoryItem>>>,
    runtime_rx: Arc<Mutex<mpsc::Receiver<MessageData>>>,
    runtime_tx: mpsc::Sender<MessageData>,
    runtime_worker_handle: Mutex<Option<JoinHandle<()>>>,

    /// Stores instances for RPC calls (as `Any` for down-casting)
    instance_registry: Arc<RwLock<HashMap<Ulid, Arc<dyn Any + Send + Sync>>>>,

    /// Registry that maps remote `RuntimeId`s to their message channels
    #[cfg(any(feature = "client", feature = "server"))]
    runtime_registry: Arc<RwLock<HashMap<RuntimeId, mpsc::Sender<Message>>>>,
}

#[crate::async_trait]
impl RuntimeTrait for Runtime {
    /// Creates a new runtime with the specified ID.
    ///
    /// The `Runtime` maintains a registry of message executors and an internal channel for
    /// scheduling message processing.
    async fn create(id: RuntimeId) -> Result<Self> {
        info!(runtime_id = %id, "Creating new RPC runtime");
        let (runtime_tx, runtime_rx) = mpsc::channel(1024);
        let rt = Self::new(RuntimeImpl {
            id,
            executors: Arc::new(RwLock::new(HashMap::new())),
            message_handlers: Arc::new(Mutex::new(HashMap::new())),
            inventory: Arc::new(Mutex::new(HashMap::new())),
            runtime_rx: Arc::new(Mutex::new(runtime_rx)),
            runtime_tx,
            runtime_worker_handle: Mutex::new(None),
            instance_registry: Arc::new(RwLock::new(HashMap::new())),

            #[cfg(any(feature = "client", feature = "server"))]
            runtime_registry: Arc::new(RwLock::new(HashMap::new())),
        });

        rt.register_handlers().await?;
        rt.runtime_worker().await;

        Ok(rt)
    }

    /// Executes a user-defined future within an RPC server context.
    ///
    /// Spawns the provided future `f` within a task-local RPC context for the given `target_id`.
    /// Incoming RPC requests for this `target_id` are dispatched to registered handlers. If an
    /// internal error is reported, the execution is aborted.
    async fn execute<F>(&self, target_id: RuntimeId, f: F) -> Result<()>
    where
        F: Future<Output = Result<()>> + Send + 'static,
    {
        let (error_tx, mut error_rx) = mpsc::channel(1);

        // Create execution context for the RPC call
        let ctx = RuntimeContext {
            target_id: target_id.clone(),
            runtime: self.clone(),
            error_tx,
        };

        let task = spawn(CONTEXT.scope(ctx, f));
        if let Some(error) = error_rx.recv().await {
            task.abort();
            Err(error)
        } else {
            task.await?
        }
    }

    /// Convenience wrapper around [`execute`] that targets the current runtime itself.
    async fn execute_local<F>(&self, f: F) -> Result<()>
    where
        F: Future<Output = Result<()>> + Send + 'static,
    {
        self.execute(self.id.clone(), f).await
    }

    /// Registers a message handler type with the runtime.
    ///
    /// The handler `H` must implement [`MessageHandler`]. After registration, the runtime will
    /// deserialize incoming messages of `H::Input` type, invoke `H::handle`, and serialize the
    /// resulting `H::Output`.
    async fn register_handler<H: MessageHandler>(&self) -> Result<()> {
        // Determine the input type id for this handler
        let type_id = <<H as MessageHandler>::Input as TypeInfo>::type_id();
        self.inventory.lock().await.insert(type_id, H::type_info());

        // Build executor that deserialises, invokes the handler, and serialises the response
        let executor = {
            let runtime = self.clone();
            let exec_fn = move |msg: Message| -> MessageFuture {
                let runtime = runtime.clone();
                Box::pin(async move {
                    // Establish a task-local RPC context for nested calls within the handler
                    let (error_tx, _error_rx) = mpsc::channel(1);
                    let ctx = RuntimeContext {
                        target_id: runtime.id.clone(),
                        runtime: runtime.clone(),
                        error_tx,
                    };

                    CONTEXT
                        .scope(ctx, async move {
                            // Deserialize input arguments
                            let reply_data: Result<_, String> = match msg.data {
                                MessageBody::Handshake(_msg) => Ok(vec![]),
                                MessageBody::Rpc(message) => match message.data {
                                    Ok(vec) => match serde_binary::from_vec::<H::Input>(
                                        vec,
                                        Endian::Little,
                                    ) {
                                        Ok(input) => match H::handle(input).await {
                                            Ok(output) => {
                                                serde_binary::to_vec(&output, Endian::Little)
                                                    .map_err(|e| e.to_string())
                                            }
                                            Err(err_str) => Err(err_str),
                                        },
                                        Err(err) => Err(err.to_string()),
                                    },
                                    Err(err_str) => Err(err_str),
                                },
                            };

                            // Build response message
                            Message {
                                target_id: msg.source_id.clone(),
                                source_id: runtime.id.clone(),
                                message_id: msg.message_id,
                                is_answer: true,
                                is_closed: false,
                                data: MessageBody::Rpc(RpcMessage {
                                    r#type: <<H as MessageHandler>::Output as TypeInfo>::type_id(),
                                    data: reply_data,
                                }),
                            }
                        })
                        .await
                })
            };
            Arc::new(Box::new(exec_fn) as Executor)
        };

        // Store executor for this input type and log registration
        info!(
            input_type = %H::Input::type_name(),
            "Registering RPC handler"
        );
        self.executors.write().await.insert(type_id, executor);
        Ok(())
    }

    /// Registers an object that implements [`Identity`] in the runtime-wide registry and returns an
    /// [`ObjectRef`] that can be shared across runtimes.
    async fn register_instance<T>(&self, instance: T) -> ObjectRef<T>
    where
        T: Identity + Any + Send + Sync + 'static,
    {
        let id = *instance.id();
        self.instance_registry
            .write()
            .await
            .insert(id, Arc::new(RwLock::new(instance)));
        ObjectRef::create(self.id.clone(), id)
    }

    /// Attempts to retrieve a previously registered instance by its [`InstanceId`].
    ///
    /// Returns `Ok` with a shared [`Arc<RwLock<T>>`] if the instance exists and can be
    /// down‑cast to the requested concrete type; otherwise returns an `Err` explaining
    /// what went wrong.
    async fn get_instance<T>(&self, id: InstanceId) -> Result<Arc<RwLock<T>>, String>
    where
        T: Identity + Any + Send + Sync + 'static,
    {
        // Clone the entry so we do not hold the registry lock across an `.await`.
        let maybe_instance = { self.instance_registry.read().await.get(&id).cloned() };

        match maybe_instance {
            Some(instance) => match instance.downcast::<RwLock<T>>() {
                Ok(typed_instance) => Ok(typed_instance),
                Err(_) => Err("type mismatch".into()),
            },
            None => Err("instance not found".into()),
        }
    }

    /// Attempts to take a previously registered instance by its [`InstanceId`].
    /// Returns `Some` if the instance exists and can be down-cast to the requested concrete type.
    /// Removes the instance from the registry.
    async fn take_instance<T>(&self, id: InstanceId) -> Option<T>
    where
        T: Identity + Any + Send + Sync + 'static,
    {
        let maybe_instance = {
            self.instance_registry
                .write()
                .await
                .remove(&id)
                .and_then(|arc_any| arc_any.downcast::<RwLock<T>>().ok())
        };

        maybe_instance.and_then(|instance| match Arc::try_unwrap(instance) {
            Ok(instance) => Some(instance.into_inner()),
            Err(_) => None,
        })
    }

    /// Sends an RPC request and awaits its response.
    ///
    /// Serializes the provided `args` into a message, sends it to the target runtime identified by
    /// `target_id`, and deserializes the response into `T::Output`.
    async fn call<T: MessageHandler>(
        &self,
        target_id: RuntimeId,
        args: T::Input,
    ) -> Result<T::Output> {
        let message_id = MessageId::new();
        trace!(
            message_id = %message_id,
            input_type = %T::Input::type_name(),
            "Sending RPC request"
        );

        let result = async move {
            // Serialize arguments and build message
            let data = Message {
                target_id,
                source_id: self.id.clone(),
                message_id,
                is_answer: false,
                is_closed: false,
                data: MessageBody::Rpc(RpcMessage {
                    r#type: <<T as MessageHandler>::Input as TypeInfo>::type_id(),
                    data: Ok(serde_binary::to_vec(&args, Endian::Little)?),
                }),
            };

            // Channel to receive answers from the runtime
            let (return_tx, mut return_rx) = mpsc::channel(1);

            // Track the handler channel
            {
                self.message_handlers
                    .lock()
                    .await
                    .insert(message_id, return_tx.clone());
            }

            // Send message to runtime for further processing
            self.runtime_tx
                .send(MessageData {
                    return_tx,
                    message: data,
                })
                .await
                .map_err(|e| anyhow!(e))?;

            // Await the response
            let ret_msg = return_rx.recv().await.context("no answer")?;

            // Deserialize the response payload
            match ret_msg.data {
                MessageBody::Rpc(message) => serde_binary::from_vec::<T::Output>(
                    message.data.map_err(|e| anyhow!(e))?,
                    Endian::Little,
                )
                .context("failed to deserialize response"),
                _ => Err(anyhow!("unsupported answer type")),
            }
        }
        .await;

        // Clean up the handler registry entry
        self.message_handlers.lock().await.remove(&message_id);

        result
    }

    #[cfg(any(feature = "client", feature = "server"))]
    /// Returns a vec of [`RuntimeId`] of connected runtimes
    async fn connected_runtimes(&self) -> Vec<RuntimeInfo> {
        self.runtime_registry
            .read()
            .await
            .keys()
            .map(|id| RuntimeInfo { id: id.clone() })
            .collect()
    }

    /// Returns a vec of all registered rpc handlers
    async fn inventory(&self) -> Vec<InventoryItem> {
        self.inventory.lock().await.values().cloned().collect()
    }

    #[cfg(feature = "client")]
    async fn connect(&self, addr: String) -> Result<JoinHandle<()>> {
        let (conn, _) = connect_async(&addr).await?;
        let (mut conn_tx, conn_rx) = conn.split();
        let (return_tx, return_rx) = mpsc::channel(16);

        self.runtime_registry
            .write()
            .await
            .insert("server".into(), return_tx.clone());

        // Send handshake
        conn_tx
            .send(tungstenite::Message::Binary(
                serde_binary::to_vec(
                    &Message {
                        target_id: "server".into(),
                        source_id: self.id.clone(),
                        message_id: Ulid::new(),
                        is_answer: false,
                        is_closed: true,
                        data: MessageBody::Handshake(HandshakeMessage {
                            runtime_id: self.id.clone(),
                        }),
                    },
                    Endian::Little,
                )?
                .into(),
            ))
            .await?;

        let runtime = self.clone();

        Ok(spawn(async move {
            // Encode all messages on the return channel and send them to the WebSocket
            let sender_handle = spawn(Self::message_encoder(conn_tx, return_rx));

            // Decode all messages from the WebSocket and pass them to the runtime
            runtime.message_decoder(return_tx, conn_rx).await;

            sender_handle.abort();
        }))
    }

    #[cfg(feature = "server")]
    async fn start_server(&self, addr: String) -> JoinHandle<Result<()>> {
        let runtime = self.clone();
        spawn(async move {
            let server = TcpListener::bind(addr).await?;

            loop {
                let (socket, _) = server.accept().await?;

                let runtime = runtime.clone();
                spawn(async move {
                    if let Err(e) = runtime.accept_connection(socket).await {
                        error!("connection error: {}", e);
                    }
                });
            }
        })
    }
}

#[crate::async_trait]
impl RuntimeInternalTrait for Runtime {
    /// Registers every RPC handler collected through the `inventory` crate.
    async fn register_handlers(&self) -> Result<()> {
        for reg in inventory::iter::<HandlerRegistration> {
            (reg.register)(self.clone()).await?;
        }
        Ok(())
    }

    /// Starts the runtime event loop.
    ///
    /// Spawns a background task that processes incoming messages.
    /// Returns a [`JoinHandle`] for the spawned worker.
    async fn runtime_worker(&self) {
        let runtime = self.clone();
        info!(runtime_id = %runtime.id, "Starting RPC runtime event loop");

        let handle = spawn(async move {
            // Continuously receive incoming messages
            let mut rx = runtime.runtime_rx.lock().await;
            while let Some(MessageData { return_tx, message }) = rx.recv().await {
                let msg_id = message.message_id;
                trace!(message_id = %msg_id, is_answer = %message.is_answer, "Received message");

                // Is the message meant for us?
                if message.target_id == runtime.id {
                    // Is it an answer to one of our previous requests?
                    if message.is_answer {
                        if let Some(handler) = {
                            let mut handlers = runtime.message_handlers.lock().await;
                            handlers.remove(&msg_id)
                        } {
                            trace!(message_id = %msg_id, "Received answer");
                            let _ = handler.send(message).await;
                        }
                    } else {
                        match message.data {
                            // Handle handshake messages (client/server builds only)
                            #[cfg(any(feature = "client", feature = "server"))]
                            MessageBody::Handshake(ref msg) => {
                                info!(runtime = %msg.runtime_id, "Connected");
                                runtime
                                    .runtime_registry
                                    .write()
                                    .await
                                    .insert(msg.runtime_id.clone(), return_tx);
                            }
                            // Handshake is a no-op in local-only builds
                            #[cfg(not(any(feature = "client", feature = "server")))]
                            MessageBody::Handshake(_) => {}
                            MessageBody::Rpc(ref msg) => {
                                if let Some(exec) =
                                    { runtime.executors.read().await.get(&msg.r#type).cloned() }
                                {
                                    trace!(message_id = %msg_id, "Dispatching to handler");
                                    let handler_exec = exec.clone();
                                    let handler_msg = message;
                                    let handler_tx = return_tx.clone();
                                    spawn(async move {
                                        let response = handler_exec(handler_msg).await;
                                        trace!(message_id = %msg_id, "Handler execution complete");
                                        if let Err(err) = handler_tx.send(response).await {
                                            error!(
                                                message_id = %msg_id,
                                                error = ?err,
                                                "Failed to send RPC response"
                                            );
                                        }
                                    });
                                } else {
                                    error!(
                                        message_id = %msg_id,
                                        message_type = ?msg.r#type,
                                        "No executor registered for message type"
                                    );
                                    let _ = return_tx
                                        .send(Message {
                                            target_id: message.source_id,
                                            source_id: runtime.id.clone(),
                                            message_id: Ulid::new(),
                                            is_answer: true,
                                            is_closed: true,
                                            data: MessageBody::Rpc(RpcMessage {
                                                r#type: msg.r#type,
                                                data: Err("no handler registered".into()),
                                            }),
                                        })
                                        .await;
                                }
                            }
                        }
                    }
                } else {
                    // Forward messages for other runtimes (client/server builds only)
                    #[cfg(any(feature = "client", feature = "server"))]
                    {
                        let target = message.target_id.clone();
                        let tx_opt =
                            { runtime.runtime_registry.read().await.get(&target).cloned() };

                        if let Some(tx) = tx_opt {
                            if let Err(err) = tx.send(message).await {
                                error!(
                                    message_id = %msg_id,
                                    target_id = %target,
                                    error = ?err,
                                    "Failed to forward message to runtime"
                                );
                            }
                        } else {
                            let server_tx_opt =
                                { runtime.runtime_registry.read().await.get("server").cloned() };
                            if let Some(tx) = server_tx_opt {
                                if let Err(err) = tx.send(message).await {
                                    error!(
                                        message_id = %msg_id,
                                        target_id = %target,
                                        error = ?err,
                                        "Failed to forward message to runtime"
                                    );
                                }
                            }
                        }
                    }
                    // In local-only builds, there should be no messages for other runtimes
                    #[cfg(not(any(feature = "client", feature = "server")))]
                    unreachable!("Received message for unknown target {}", message.target_id);
                }
            }
        });

        self.runtime_worker_handle.lock().await.replace(handle);
    }

    #[cfg(any(feature = "client", feature = "server"))]
    async fn message_encoder<S: Sink<tungstenite::Message> + Unpin + Send>(
        mut tx: S,
        mut rx: mpsc::Receiver<Message>,
    ) where
        <S as Sink<tungstenite::Message>>::Error: std::fmt::Display,
    {
        while let Some(message) = rx.recv().await {
            match serde_binary::to_vec(&message, Endian::Little) {
                Ok(message) => {
                    let message = tungstenite::Message::Binary(message.into());
                    if let Err(e) = tx.send(message).await {
                        error!("error sending message: {}", e);
                    }
                }
                Err(e) => error!("error sending message: {}", e),
            }
        }
    }

    #[cfg(any(feature = "client", feature = "server"))]
    async fn message_decoder<
        S: Stream<Item = std::result::Result<tungstenite::Message, tungstenite::Error>> + Unpin + Send,
    >(
        &self,
        tx: mpsc::Sender<Message>,
        mut rx: S,
    ) {
        while let Some(message) = rx.next().await {
            match message {
                Ok(tungstenite::Message::Binary(message)) => {
                    match serde_binary::from_slice(&message, Endian::Little) {
                        Ok(message) => {
                            let _ = self
                                .runtime_tx
                                .send(MessageData {
                                    message,
                                    return_tx: tx.clone(),
                                })
                                .await;
                        }
                        Err(e) => error!("error decoding message: {}", e),
                    }
                }
                Ok(_) => {
                    error!("unsupported message type");
                    break;
                }
                Err(e) => {
                    error!("connection error: {}", e);
                    break;
                }
            }
        }
    }

    #[cfg(feature = "server")]
    async fn accept_connection(&self, stream: TcpStream) -> Result<()> {
        let ws = tokio_tungstenite::accept_async(stream).await?;
        let (stream_tx, stream_rx) = ws.split();
        let (return_tx, return_rx) = mpsc::channel(16);

        let transfer_handle = spawn(Self::message_encoder(stream_tx, return_rx));
        self.message_decoder(return_tx, stream_rx).await;
        transfer_handle.abort();

        Ok(())
    }
}