clawft_kernel/workload_ctl/
transport.rs

1//! How the control plane reaches a `workload-host`, and how a node serves
2//! one: mesh TCP (optionally Noise XX, as the kernel mesh listener does)
3//! and in-process (same wire, in-memory streams) for local and Seed-adapter
4//! nodes and for hermetic tests.
5
6use std::collections::HashMap;
7use std::sync::{Arc, RwLock};
8
9use async_trait::async_trait;
10
11use crate::mesh::{MeshError, MeshStream, MeshTransport, TransportListener};
12use crate::mesh_noise::{NoiseChannel, NoiseConfig};
13use crate::mesh_tcp::TcpTransport;
14use crate::mesh_test_support::connected_pair;
15
16use super::host_service::WorkloadHostService;
17use super::session::{NoiseStream, serve_connection};
18
19/// Address scheme for in-process hosts.
20pub const MEM_SCHEME: &str = "mem://";
21
22/// Opens a stream to a target address.
23#[async_trait]
24pub trait CtlConnector: Send + Sync {
25    /// Connect to `addr`.
26    async fn connect(&self, addr: &str) -> Result<Box<dyn MeshStream>, MeshError>;
27}
28
29fn noise_config() -> NoiseConfig {
30    NoiseConfig {
31        pattern: crate::mesh_noise::NoisePattern::XX,
32        local_private_key: rand::random(),
33        remote_static_key: None,
34    }
35}
36
37/// Mesh TCP, with Noise XX when `noise` is set; `mem://` addresses go to
38/// registered in-process hosts.
39#[derive(Default)]
40pub struct MeshConnector {
41    noise: bool,
42    local: RwLock<HashMap<String, Arc<WorkloadHostService>>>,
43}
44
45impl MeshConnector {
46    /// Connector; `noise` wraps TCP streams in Noise XX.
47    pub fn new(noise: bool) -> Self {
48        Self {
49            noise,
50            local: RwLock::new(HashMap::new()),
51        }
52    }
53
54    /// Register an in-process host at `mem://<name>` and return the address.
55    /// Used for this node's own host and for Seed adapters on an
56    /// operator-assigned node id.
57    pub fn register_local(&self, name: &str, svc: Arc<WorkloadHostService>) -> String {
58        let addr = format!("{MEM_SCHEME}{name}");
59        if let Ok(mut m) = self.local.write() {
60            m.insert(addr.clone(), svc);
61        }
62        addr
63    }
64}
65
66#[async_trait]
67impl CtlConnector for MeshConnector {
68    async fn connect(&self, addr: &str) -> Result<Box<dyn MeshStream>, MeshError> {
69        if addr.starts_with(MEM_SCHEME) {
70            let svc = self
71                .local
72                .read()
73                .ok()
74                .and_then(|m| m.get(addr).cloned())
75                .ok_or_else(|| MeshError::PeerNotConnected(addr.to_string()))?;
76            let (client, server) = connected_pair().await?;
77            tokio::spawn(async move {
78                if let Err(e) = serve_connection(Box::new(server), svc).await {
79                    tracing::debug!(error = %e, "in-process workload-host session ended");
80                }
81            });
82            return Ok(Box::new(client));
83        }
84        let stream = TcpTransport.connect(addr).await?;
85        if !self.noise {
86            return Ok(stream);
87        }
88        let ch = NoiseChannel::initiate(stream, &noise_config()).await?;
89        Ok(Box::new(NoiseStream::new(Box::new(ch))))
90    }
91}
92
93/// Most concurrent `workload-host` sessions a listener serves.
94pub const MAX_SESSIONS: usize = 32;
95
96/// Accept connections on `listener` and serve each as a `workload-host`
97/// session (Noise XX responder when `noise` is set). Runs until the
98/// listener fails. At most [`MAX_SESSIONS`] at once; more are closed.
99pub async fn serve_listener(
100    listener: Box<dyn TransportListener>,
101    svc: Arc<WorkloadHostService>,
102    noise: bool,
103) -> Result<(), MeshError> {
104    serve_listener_with(listener, svc, noise, MAX_SESSIONS).await
105}
106
107/// [`serve_listener`] with an explicit session cap.
108pub async fn serve_listener_with(
109    mut listener: Box<dyn TransportListener>,
110    svc: Arc<WorkloadHostService>,
111    noise: bool,
112    max_sessions: usize,
113) -> Result<(), MeshError> {
114    let slots = Arc::new(tokio::sync::Semaphore::new(max_sessions));
115    loop {
116        let (mut stream, peer) = listener.accept().await?;
117        let Ok(slot) = slots.clone().try_acquire_owned() else {
118            tracing::debug!(%peer, "workload-host session limit reached; connection closed");
119            let _ = stream.close().await;
120            continue;
121        };
122        let svc = svc.clone();
123        tokio::spawn(async move {
124            let _slot = slot;
125            let stream: Box<dyn MeshStream> = if noise {
126                match NoiseChannel::respond(stream, &noise_config()).await {
127                    Ok(ch) => Box::new(NoiseStream::new(Box::new(ch))),
128                    Err(e) => {
129                        tracing::warn!(%peer, error = %e, "workload-host noise handshake failed");
130                        return;
131                    }
132                }
133            } else {
134                stream
135            };
136            if let Err(e) = serve_connection(stream, svc).await {
137                tracing::debug!(%peer, error = %e, "workload-host session ended");
138            }
139        });
140    }
141}
142
143/// Bind a TCP listener for [`serve_listener`].
144pub async fn listen_tcp(addr: &str) -> Result<Box<dyn TransportListener>, MeshError> {
145    TcpTransport.listen(addr).await
146}
clawft_kernel/workload_ctl/transport.rs

clawft_kernel/workload_ctl/
transport.rs
