ractor_cluster 0.16.4

// Copyright (c) Sean Lawlor
//
// This source code is licensed under both the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree.

//! A [RemoteActor] is an actor which handles serialized messages, and represents an actor
//! running on a remote `node()` server. See [crate::node::NodeServer] for more on inter-node
//! protocols

use std::collections::BTreeMap;
use std::ops::Bound::{Excluded, Unbounded};

use ractor::cast;
use ractor::concurrency::JoinHandle;
use ractor::message::SerializedMessage;
use ractor::Actor;
use ractor::ActorCell;
use ractor::ActorId;
use ractor::ActorName;
use ractor::ActorProcessingErr;
use ractor::ActorRef;
use ractor::RpcReplyPort;
use ractor::SpawnErr;
use ractor_cluster_derive::RactorMessage;

use crate::node::NodeSessionMessage;
use crate::NodeId;

#[cfg(test)]
mod tests;

/// Maximum number of pending requests inspected when handling a message.
const PENDING_REQUEST_CLEANUP_BUDGET: usize = 16;

/// A [RemoteActor] is an actor which represents an actor on another node
///
/// A [RemoteActor] handles serialized messages without decoding them, forwarding them
/// to the remote system for decoding + handling by the real implementation.
/// Therefore [RemoteActor]s can be thought of as a "shim" to a real actor on a remote system
pub(crate) struct RemoteActor;

impl RemoteActor {
    /// Spawn a [RemoteActor] with a supervisor (which should always be a [super::node::NodeSession])
    /// and automatically start the actor
    ///
    /// * `name`: A name to give the actor. Useful for global referencing or debug printing
    /// * `pid`: The actor's local id on the remote system
    /// * `node_id` The id of the [super::node::NodeSession]. Alongside `pid` this makes for a unique actor identifier
    /// * `supervisor`: The [super::node::NodeSession]'s [ActorCell] handle which will be linked in
    ///   the supervision tree
    ///
    /// Returns a [Ok((ActorRef, JoinHandle<()>))] upon successful start, denoting the actor reference
    /// along with the join handle which will complete when the actor terminates. Returns [Err(SpawnErr)] if
    /// the actor failed to start
    pub(crate) async fn spawn_linked(
        self,
        session: ActorRef<super::node::NodeSessionMessage>,
        name: Option<ActorName>,
        pid: u64,
        node_id: NodeId,
        supervisor: ActorCell,
    ) -> Result<(ActorRef<RemoteActorMessage>, JoinHandle<()>), SpawnErr> {
        let actor_id = ActorId::Remote { node_id, pid };
        ractor::ActorRuntime::<Self>::spawn_linked_remote(name, self, actor_id, session, supervisor)
            .await
    }
}

pub(crate) struct RemoteActorState {
    message_tag: u64,
    /// The map of <message_tag, serialized_rpc_reply> port for network
    /// handling of [SerializedMessage::CallReply]s
    pending_requests: BTreeMap<u64, RpcReplyPort<Vec<u8>>>,
    /// Last request inspected for cancellation.
    pending_request_cleanup_cursor: Option<u64>,

    /// Owning session
    session: ActorRef<crate::node::NodeSessionMessage>,
}

impl RemoteActorState {
    fn new(session: ActorRef<crate::node::NodeSessionMessage>) -> Self {
        Self {
            session,
            message_tag: 0,
            pending_requests: BTreeMap::new(),
            pending_request_cleanup_cursor: None,
        }
    }

    /// Increment the message tag and return the "next" value
    fn get_and_increment_mtag(&mut self) -> u64 {
        self.message_tag += 1;
        self.message_tag
    }

    /// Incrementally reclaim reply ports whose receivers were cancelled or timed out.
    fn cleanup_closed_pending_requests(&mut self) {
        let cleanup_count = self
            .pending_requests
            .len()
            .min(PENDING_REQUEST_CLEANUP_BUDGET);
        if cleanup_count == 0 {
            self.pending_request_cleanup_cursor = None;
            return;
        }

        let mut tags = [0; PENDING_REQUEST_CLEANUP_BUDGET];
        let mut tag_count = 0;

        if let Some(cursor) = self.pending_request_cleanup_cursor {
            for (&tag, _) in self
                .pending_requests
                .range((Excluded(cursor), Unbounded))
                .take(cleanup_count)
            {
                tags[tag_count] = tag;
                tag_count += 1;
            }
        }

        for (&tag, _) in self.pending_requests.iter().take(cleanup_count - tag_count) {
            tags[tag_count] = tag;
            tag_count += 1;
        }

        for tag in &tags[..tag_count] {
            if self
                .pending_requests
                .get(tag)
                .is_some_and(RpcReplyPort::is_closed)
            {
                self.pending_requests.remove(tag);
            }
        }

        self.pending_request_cleanup_cursor = if self.pending_requests.is_empty() {
            None
        } else {
            tags[..tag_count].last().copied()
        };
    }

    fn remove_pending_request(&mut self, message_tag: u64) -> Option<RpcReplyPort<Vec<u8>>> {
        let port = self.pending_requests.remove(&message_tag);
        if self.pending_requests.is_empty() {
            self.pending_request_cleanup_cursor = None;
        }
        port
    }
}

// Placeholder message for a remote actor, as we won't actually
// handle anything but serialized messages on this channel
#[derive(RactorMessage)]
pub(crate) struct RemoteActorMessage;

#[cfg_attr(feature = "async-trait", ractor::async_trait)]
impl Actor for RemoteActor {
    type Msg = RemoteActorMessage;
    type State = RemoteActorState;
    type Arguments = ActorRef<crate::node::NodeSessionMessage>;
    async fn pre_start(
        &self,
        _myself: ActorRef<Self::Msg>,
        session: ActorRef<crate::node::NodeSessionMessage>,
    ) -> Result<Self::State, ActorProcessingErr> {
        Ok(Self::State::new(session))
    }

    async fn handle(
        &self,
        _myself: ActorRef<Self::Msg>,
        _message: Self::Msg,
        _state: &mut Self::State,
    ) -> Result<(), ActorProcessingErr> {
        Err(From::from("`RemoteActor`s cannot handle local messages!"))
    }

    async fn handle_serialized(
        &self,
        myself: ActorRef<Self::Msg>,
        message: SerializedMessage,
        state: &mut Self::State,
    ) -> Result<(), ActorProcessingErr> {
        state.cleanup_closed_pending_requests();

        // get the local pid on the remote system
        let to = myself.get_id().pid();
        // messages should be forwarded over the network link (i.e. sent through the node session) to the intended
        // target node's relevant actor. The receiving runtime NodeSession will decode the message and pass it up
        // to the parent. However `SerializedMessage::CallReply` is a network reply to a send call request
        match message {
            SerializedMessage::Call {
                args,
                reply,
                variant,
                metadata,
            } => {
                // Handle Call
                let tag = state.get_and_increment_mtag();
                let node_msg = crate::protocol::node::NodeMessage {
                    msg: Some(crate::protocol::node::node_message::Msg::Call(
                        crate::protocol::node::Call {
                            to,
                            tag,
                            what: args,
                            timeout_ms: reply.get_timeout().map(|t| t.as_millis() as u64),
                            variant,
                            metadata,
                        },
                    )),
                };
                state.pending_requests.insert(tag, reply);
                if cast!(state.session, NodeSessionMessage::SendMessage(node_msg)).is_err() {
                    state.remove_pending_request(tag);
                }
            }
            SerializedMessage::Cast {
                args,
                variant,
                metadata,
            } => {
                // Handle Cast
                let node_msg = crate::protocol::node::NodeMessage {
                    msg: Some(crate::protocol::node::node_message::Msg::Cast(
                        crate::protocol::node::Cast {
                            to,
                            what: args,
                            variant,
                            metadata,
                        },
                    )),
                };
                let _ = cast!(state.session, NodeSessionMessage::SendMessage(node_msg));
            }
            SerializedMessage::CallReply(message_tag, reply_data) => {
                // Handle the reply to a "Call" message
                if let Some(port) = state.remove_pending_request(message_tag) {
                    let _ = port.send(reply_data);
                }
            }
        }
        Ok(())
    }
}