Skip to main content

netget/server/statsd/
mod.rs

1//! StatsD collector with bounded, sequential batch handling and no reply traffic.
2pub mod actions;
3pub mod codec;
4
5use crate::protocol::{Event, SpawnContext};
6use crate::server::connection::ConnectionId;
7use crate::state::{
8    server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo},
9    AccessLogOwner,
10};
11use crate::{console_debug, console_error, console_info};
12use anyhow::Result;
13use codec::{Dialect, DEFAULT_DIALECT, DEFAULT_LLM_FALLBACK, MAX_DATAGRAM_BYTES};
14use serde_json::json;
15use tokio::net::UdpSocket;
16
17pub struct StatsdServer;
18impl StatsdServer {
19    pub async fn spawn(ctx: SpawnContext) -> Result<std::net::SocketAddr> {
20        let dialect = Dialect::parse(
21            &ctx.startup_params
22                .as_ref()
23                .map(|p| p.get_optional_string("dialect"))
24                .transpose()?
25                .flatten()
26                .unwrap_or_else(|| DEFAULT_DIALECT.into()),
27        )?;
28        let llm_fallback = ctx
29            .startup_params
30            .as_ref()
31            .map(|p| p.get_optional_bool("llm_fallback"))
32            .transpose()?
33            .flatten()
34            .unwrap_or(DEFAULT_LLM_FALLBACK);
35        let socket = UdpSocket::bind(ctx.legacy_listen_addr()).await?;
36        let local_addr = socket.local_addr()?;
37        console_info!(
38            ctx.status_tx,
39            "StatsD listening on {} ({}, llm_fallback={})",
40            local_addr,
41            dialect.as_str(),
42            llm_fallback
43        );
44        let registrar = ctx.state.clone();
45        let server_id = ctx.server_id;
46        let handle = tokio::spawn(async move {
47            // +1 detects a truncated datagram: anything above the cap is rejected,
48            // even if its prefix would be a valid message. No unbounded task queue.
49            let mut buffer = vec![0u8; MAX_DATAGRAM_BYTES + 1];
50            let protocol = actions::StatsdProtocol::new();
51            loop {
52                let (n, peer) = match socket.recv_from(&mut buffer).await {
53                    Ok(received) => received,
54                    Err(error) => {
55                        console_error!(ctx.status_tx, "StatsD receive failed: {}", error);
56                        break;
57                    }
58                };
59                let records = match codec::parse_datagram(&buffer[..n], dialect) {
60                    Ok(records) => records,
61                    Err(error) => {
62                        console_debug!(
63                            ctx.status_tx,
64                            "StatsD rejected datagram from {}: {}",
65                            peer,
66                            error
67                        );
68                        ctx.state
69                            .record_access_log(
70                                AccessLogOwner::Server(server_id.as_u32()),
71                                "StatsD",
72                                None,
73                                "statsd_invalid_datagram",
74                                json!({"source_addr":peer.to_string(),"received_bytes":n}),
75                                vec![json!({"error":error.to_string()})],
76                            )
77                            .await;
78                        continue;
79                    }
80                };
81                // Only the batch currently being handled needs a peer row. Remove it
82                // afterward rather than accumulate one connection per telemetry packet.
83                let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
84                let now = crate::utils::clock::Instant::now();
85                ctx.state
86                    .add_connection_to_server(
87                        server_id,
88                        ConnectionState {
89                            id,
90                            remote_addr: peer,
91                            local_addr,
92                            bytes_sent: 0,
93                            bytes_received: n as u64,
94                            packets_sent: 0,
95                            packets_received: 1,
96                            last_activity: now,
97                            status: ConnectionStatus::Active,
98                            status_changed_at: now,
99                            protocol_info: ProtocolConnectionInfo::empty(),
100                        },
101                    )
102                    .await;
103                let event = Event::new(
104                    &actions::STATSD_BATCH_EVENT,
105                    json!({"records":records,"record_count":records.len(),"source_addr":peer.to_string(),"dialect":dialect.as_str()}),
106                );
107                let configured = ctx
108                    .state
109                    .get_event_handler_config(server_id)
110                    .await
111                    .is_some_and(|config| config.find_handler("statsd_batch").is_some());
112                if llm_fallback || configured {
113                    match crate::llm::action_helper::call_llm(
114                        &ctx.llm_client,
115                        &ctx.state,
116                        server_id,
117                        Some(id),
118                        &event,
119                        &protocol,
120                    )
121                    .await
122                    {
123                        Ok(result) => {
124                            for message in result.messages {
125                                console_info!(ctx.status_tx, "{}", message);
126                            }
127                        }
128                        Err(error) => {
129                            console_error!(
130                                ctx.status_tx,
131                                "StatsD decision=fail_closed_handler_error from {}: {}",
132                                peer,
133                                error
134                            );
135                        }
136                    }
137                } else {
138                    ctx.state
139                        .record_access_log(
140                            AccessLogOwner::Server(server_id.as_u32()),
141                            "StatsD",
142                            Some(id.as_u32()),
143                            "statsd_batch",
144                            event.data,
145                            vec![json!({"type":"collect_statsd_batch"})],
146                        )
147                        .await;
148                }
149                ctx.state.remove_connection_from_server(server_id, id).await;
150                let _ = ctx.status_tx.send("__UPDATE_UI__".into());
151            }
152        });
153        registrar.register_server_task(server_id, handle).await;
154        Ok(local_addr)
155    }
156}