netget/server/statsd/
mod.rs1pub 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 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 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}