1pub mod actions;
3pub mod codec;
4use crate::protocol::{Event, SpawnContext};
5use crate::server::{
6 accept_bounded::{accept_bounded, ConnectionLimiter, DEFAULT_MAX_CONNECTIONS},
7 connection::ConnectionId,
8};
9use crate::state::{
10 server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo},
11 AccessLogOwner,
12};
13use crate::{console_error, console_info};
14use anyhow::{Context, Result};
15use serde_json::json;
16use std::{
17 net::SocketAddr,
18 time::{Duration, SystemTime, UNIX_EPOCH},
19};
20use tokio::{
21 io::AsyncReadExt,
22 net::{TcpListener, TcpStream},
23};
24
25pub const READ_TIMEOUT: Duration = Duration::from_secs(30);
26pub struct GraphiteServer;
27impl GraphiteServer {
28 pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
29 let llm_fallback = ctx
30 .startup_params
31 .as_ref()
32 .map(|p| p.get_optional_bool("llm_fallback"))
33 .transpose()?
34 .flatten()
35 .unwrap_or(codec::DEFAULT_LLM_FALLBACK);
36 let listener = TcpListener::bind(ctx.legacy_listen_addr()).await?;
37 let local = listener.local_addr()?;
38 console_info!(
39 ctx.status_tx,
40 "Graphite Carbon TCP listening on {} (llm_fallback={})",
41 local,
42 llm_fallback
43 );
44 let registrar = ctx.state.clone();
45 let server_id = ctx.server_id;
46 let task = tokio::spawn(async move {
47 let limiter = ConnectionLimiter::new(DEFAULT_MAX_CONNECTIONS);
48 loop {
49 let (socket, peer, permit) = match accept_bounded(
51 &listener,
52 &limiter,
53 b"",
54 "Graphite",
55 Some(&ctx.status_tx),
56 )
57 .await
58 {
59 Ok(accepted) => accepted,
60 Err(error) => {
61 console_error!(ctx.status_tx, "Graphite accept failed: {}", error);
62 break;
63 }
64 };
65 let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
66 let now = crate::utils::clock::Instant::now();
67 ctx.state
68 .add_connection_to_server(
69 server_id,
70 ConnectionState {
71 id,
72 remote_addr: peer,
73 local_addr: local,
74 bytes_sent: 0,
75 bytes_received: 0,
76 packets_sent: 0,
77 packets_received: 0,
78 last_activity: now,
79 status: ConnectionStatus::Active,
80 status_changed_at: now,
81 protocol_info: ProtocolConnectionInfo::empty(),
82 },
83 )
84 .await;
85 let child_ctx = ctx.clone();
86 let child = tokio::spawn(async move {
87 let _permit = permit;
88 if let Err(error) = session(socket, peer, id, &child_ctx, llm_fallback).await {
89 console_error!(
90 child_ctx.status_tx,
91 "Graphite peer {} closed: {}",
92 peer,
93 error
94 );
95 child_ctx
96 .state
97 .record_access_log(
98 AccessLogOwner::Server(server_id.as_u32()),
99 "Graphite",
100 Some(id.as_u32()),
101 "graphite_invalid_stream",
102 json!({"source_addr":peer.to_string()}),
103 vec![json!({"error":error.to_string()})],
104 )
105 .await;
106 }
107 child_ctx
108 .state
109 .remove_connection_from_server(server_id, id)
110 .await;
111 let _ = child_ctx.status_tx.send("__UPDATE_UI__".into());
112 });
113 ctx.state.register_server_task(server_id, child).await;
114 }
115 });
116 registrar.register_server_task(server_id, task).await;
117 Ok(local)
118 }
119}
120async fn session(
121 mut socket: TcpStream,
122 peer: SocketAddr,
123 id: ConnectionId,
124 ctx: &SpawnContext,
125 llm_fallback: bool,
126) -> Result<()> {
127 let mut decoder = codec::Decoder::default();
128 let mut buffer = [0u8; codec::READ_BYTES];
129 let protocol = actions::GraphiteProtocol::new();
130 loop {
131 let deadline = tokio::time::Instant::now() + READ_TIMEOUT;
134 let mut metrics = loop {
135 let metrics = decoder.next_batch()?;
136 if !metrics.is_empty() {
137 break metrics;
138 }
139 let n = tokio::time::timeout_at(deadline, socket.read(&mut buffer))
140 .await
141 .context("Carbon metric frame read deadline exceeded")??;
142 if n == 0 {
143 return decoder.finish();
144 }
145 ctx.state
146 .update_connection_stats(ctx.server_id, id, Some(n as u64), None, Some(1), None)
147 .await;
148 decoder.feed(&buffer[..n])?;
149 };
150 let received_time = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs_f64();
151 for metric in &mut metrics {
152 if metric.timestamp == -1.0 {
153 metric.timestamp = received_time;
154 }
155 }
156 let event = Event::new(
157 &actions::GRAPHITE_BATCH_EVENT,
158 json!({"record_count":metrics.len(),"metrics":metrics,"source_addr":peer.to_string()}),
159 );
160 let configured = ctx
161 .state
162 .get_event_handler_config(ctx.server_id)
163 .await
164 .is_some_and(|c| c.find_handler("graphite_batch").is_some());
165 if llm_fallback || configured {
166 match crate::llm::action_helper::call_llm(
167 &ctx.llm_client,
168 &ctx.state,
169 ctx.server_id,
170 Some(id),
171 &event,
172 &protocol,
173 )
174 .await
175 {
176 Ok(result) => {
177 for message in result.messages {
178 console_info!(ctx.status_tx, "{}", message);
179 }
180 }
181 Err(error) => {
182 console_error!(
183 ctx.status_tx,
184 "Graphite decision=fail_closed_handler_error: {}",
185 error
186 );
187 return Err(error);
188 }
189 }
190 } else {
191 ctx.state
192 .record_access_log(
193 AccessLogOwner::Server(ctx.server_id.as_u32()),
194 "Graphite",
195 Some(id.as_u32()),
196 "graphite_batch",
197 event.data,
198 vec![json!({"type":"collect_graphite_batch"})],
199 )
200 .await;
201 }
202 let _ = ctx.status_tx.send("__UPDATE_UI__".into());
203 }
204}