orlando-cluster 0.1.0

A virtual actor framework in Rust, inspired by Microsoft Orleans.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

use std::sync::Arc;
use std::time::Duration;

use dashmap::DashMap;

use crate::error::ClusterError;
use crate::proto::cluster_gateway_client::ClusterGatewayClient;
use crate::proto::grain_transport_client::GrainTransportClient;
use crate::proto::membership_client::MembershipClient;

/// Timeout for establishing new gRPC connections.
const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);

pub struct ConnectionPool {
    transports: DashMap<String, GrainTransportClient<tonic::transport::Channel>>,
    memberships: DashMap<String, MembershipClient<tonic::transport::Channel>>,
    gateways: DashMap<String, ClusterGatewayClient<tonic::transport::Channel>>,
    tls_config: Option<tonic::transport::ClientTlsConfig>,
    auth_token: Option<String>,
}

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

impl ConnectionPool {
    pub fn new() -> Self {
        Self {
            transports: DashMap::new(),
            memberships: DashMap::new(),
            gateways: DashMap::new(),
            tls_config: None,
            auth_token: None,
        }
    }

    pub fn with_tls(tls: tonic::transport::ClientTlsConfig) -> Self {
        Self {
            tls_config: Some(tls),
            ..Self::new()
        }
    }

    pub fn with_tls_and_auth(tls: tonic::transport::ClientTlsConfig, token: String) -> Self {
        Self {
            tls_config: Some(tls),
            auth_token: Some(token),
            ..Self::new()
        }
    }

    pub fn with_auth(token: String) -> Self {
        Self {
            auth_token: Some(token),
            ..Self::new()
        }
    }

    /// Wrap a value in a `tonic::Request` with the auth token attached (if configured).
    pub fn authorized_request<T>(&self, inner: T) -> tonic::Request<T> {
        let mut request = tonic::Request::new(inner);
        if let Some(ref token) = self.auth_token
            && let Ok(value) = token.parse()
        {
            request.metadata_mut().insert("authorization", value);
        }
        request
    }

    async fn connect_channel(&self, endpoint: &str) -> Result<tonic::transport::Channel, ClusterError> {
        let uri = if self.tls_config.is_some() {
            format!("https://{}", endpoint)
        } else {
            format!("http://{}", endpoint)
        };
        let mut ep = tonic::transport::Endpoint::from_shared(uri)
            .map_err(|e| ClusterError::Transport(e.to_string()))?
            .connect_timeout(CONNECT_TIMEOUT);

        if let Some(ref tls) = self.tls_config {
            ep = ep
                .tls_config(tls.clone())
                .map_err(|e| ClusterError::Transport(e.to_string()))?;
        }

        ep.connect()
            .await
            .map_err(|e| ClusterError::Transport(e.to_string()))
    }

    pub async fn get_transport(
        self: &Arc<Self>,
        endpoint: &str,
    ) -> Result<GrainTransportClient<tonic::transport::Channel>, ClusterError> {
        if let Some(client) = self.transports.get(endpoint) {
            return Ok(client.clone());
        }

        let channel = self.connect_channel(endpoint).await?;
        let client = GrainTransportClient::new(channel);

        self.transports
            .insert(endpoint.to_string(), client.clone());
        Ok(client)
    }

    pub async fn get_membership(
        self: &Arc<Self>,
        endpoint: &str,
    ) -> Result<MembershipClient<tonic::transport::Channel>, ClusterError> {
        if let Some(client) = self.memberships.get(endpoint) {
            return Ok(client.clone());
        }

        let channel = self.connect_channel(endpoint).await?;
        let client = MembershipClient::new(channel);

        self.memberships
            .insert(endpoint.to_string(), client.clone());
        Ok(client)
    }

    pub async fn get_gateway(
        self: &Arc<Self>,
        endpoint: &str,
    ) -> Result<ClusterGatewayClient<tonic::transport::Channel>, ClusterError> {
        if let Some(client) = self.gateways.get(endpoint) {
            return Ok(client.clone());
        }

        let channel = self.connect_channel(endpoint).await?;
        let client = ClusterGatewayClient::new(channel);

        self.gateways
            .insert(endpoint.to_string(), client.clone());
        Ok(client)
    }

    /// Remove cached connections for an endpoint.
    ///
    /// Called automatically when SWIM declares a member dead. Callers should
    /// also call this after persistent connection errors to force reconnection.
    ///
    /// Note: tonic `Channel` handles reconnection internally for transient
    /// failures — this method is for permanent removals (dead silos) or
    /// forcing a fresh connection after repeated errors.
    pub fn remove(&self, endpoint: &str) {
        self.transports.remove(endpoint);
        self.memberships.remove(endpoint);
        self.gateways.remove(endpoint);
    }
}