kapsl-ipc 0.1.0

Inter-process communication protocol for Kapsl engine clients
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

use async_trait::async_trait;
use kapsl_scheduler::ReplicaScheduler;
use kapsl_transport::{TransportError, TransportServer};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::net::TcpListener;

/// TCP-based server for network communication across different computers
pub struct TcpServer {
    bind_addr: String,
    port: u16,
    scheduler_lookup: crate::server::SchedulerLookup,
    auth_token: Option<Arc<str>>,
}

impl TcpServer {
    pub fn new(
        bind_addr: &str,
        port: u16,
        schedulers: HashMap<u32, Arc<dyn ReplicaScheduler + Send + Sync>>,
    ) -> Self {
        let schedulers = Arc::new(schedulers);
        let scheduler_lookup: crate::server::SchedulerLookup =
            Arc::new(move |model_id| schedulers.get(&model_id).cloned());
        Self::new_with_lookup(bind_addr, port, scheduler_lookup)
    }

    pub fn new_with_lookup(
        bind_addr: &str,
        port: u16,
        scheduler_lookup: crate::server::SchedulerLookup,
    ) -> Self {
        Self {
            bind_addr: bind_addr.to_string(),
            port,
            scheduler_lookup,
            auth_token: None,
        }
    }

    pub fn with_auth_token(mut self, token: impl Into<String>) -> Self {
        self.auth_token = Some(Arc::from(token.into().as_str()));
        self
    }

    async fn run_internal(&self) -> std::io::Result<()> {
        let addr = format!("{}:{}", self.bind_addr, self.port);
        let listener = TcpListener::bind(&addr).await?;
        let scheduler_lookup = self.scheduler_lookup.clone();
        let auth_token = self.auth_token.clone();

        log::info!("TCP Server listening on {}", addr);
        log::info!("TCP Server bound to {}", addr);

        loop {
            let (stream, peer_addr) = listener.accept().await?;
            let scheduler_lookup = scheduler_lookup.clone();
            let auth_token = auth_token.clone();

            log::info!("New TCP connection from {}", peer_addr);

            tokio::spawn(async move {
                if let Err(e) =
                    crate::server::handle_connection(stream, scheduler_lookup, None, auth_token)
                        .await
                {
                    log::error!("Connection error: {}", e);
                }
                log::info!("TCP connection closed from {}", peer_addr);
            });
        }
    }
}

#[async_trait]
impl TransportServer for TcpServer {
    async fn run(&self) -> Result<(), TransportError> {
        self.run_internal().await.map_err(TransportError::Io)
    }

    async fn shutdown(&self) -> Result<(), TransportError> {
        // TCP listeners don't need explicit cleanup
        Ok(())
    }

    fn transport_type(&self) -> &'static str {
        "tcp"
    }
}