salix_controller/
lib.rs

1//! Salix controller
2use std::{path::PathBuf, sync::Arc, time::Duration};
3
4use anyhow::Result;
5use clap::Parser;
6use salix_config::get_config;
7use salix_proto::{
8    MessageRequest, RegistrationRequest,
9    controller_service_server::{ControllerService, ControllerServiceServer},
10};
11use tokio::{sync::Mutex, task::JoinSet, time::sleep};
12use tonic::{Request, Response, Status, transport::Server};
13
14/// Our own lovely constant
15pub const ALPN_QUIC_SALIX: &[&[u8]] = &[b"salix"];
16
17/// CLI arguments
18#[derive(Parser)]
19#[command(version, about, long_about = None)]
20pub struct Cli {
21    #[arg(short, long, value_name = "FILE")]
22    config: Option<PathBuf>,
23}
24
25#[derive(Debug)]
26struct Agent {
27    agent_id: String,
28    hostname: String,
29    version: String,
30}
31
32#[derive(Debug, Default)]
33struct ControllerState {
34    agents: Vec<Agent>,
35}
36
37#[derive(Debug, Default)]
38struct Controller {
39    state: Arc<Mutex<ControllerState>>,
40}
41
42impl Controller {
43    fn new(state: Arc<Mutex<ControllerState>>) -> Self {
44        Self { state }
45    }
46}
47
48#[tonic::async_trait]
49impl ControllerService for Controller {
50    async fn register(
51        &self,
52        request: Request<RegistrationRequest>,
53    ) -> Result<Response<()>, Status> {
54        let request = request.into_inner();
55        if request.agent_id.is_none() {
56            return Err(Status::invalid_argument("agent_id is required"));
57        }
58
59        let mut state = self.state.lock().await;
60        state.agents.push(Agent {
61            agent_id: request.agent_id().to_owned(),
62            hostname: request.hostname().to_owned(),
63            version: request.version().to_owned(),
64        });
65
66        Ok(Response::new(()))
67    }
68
69    async fn message(&self, _request: Request<MessageRequest>) -> Result<Response<()>, Status> {
70        Ok(Response::new(()))
71    }
72}
73
74async fn print_state(state_m: Arc<Mutex<ControllerState>>) -> Result<(), tonic::transport::Error> {
75    loop {
76        let state = state_m.lock().await;
77        println!("Currently registered agents:");
78        for agent in &state.agents {
79            println!(
80                "{agent_id} {hostname} {version}",
81                agent_id = agent.agent_id,
82                hostname = agent.hostname,
83                version = agent.version
84            );
85        }
86        drop(state);
87        sleep(Duration::from_secs(10)).await;
88    }
89}
90
91/// Main function
92#[tokio::main]
93pub async fn run(cli: Cli) -> Result<()> {
94    let config = get_config(cli.config)?;
95
96    let state = Arc::new(Mutex::new(ControllerState::default()));
97
98    let controller = Controller::new(state.clone());
99
100    let server = Server::builder()
101        .add_service(ControllerServiceServer::new(controller))
102        .serve(config.controller.listen);
103    let state_printer = print_state(state.clone());
104
105    let mut set = JoinSet::new();
106    set.spawn(server);
107    set.spawn(state_printer);
108    set.join_all().await;
109
110    Ok(())
111}
salix_controller/lib.rs

salix_controller/
lib.rs
