Skip to main content

netget/server/graphite/
mod.rs

1//! One-way Carbon plaintext collector with bounded stream parsing and owned tasks.
2pub 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                // Carbon has no response grammar: refuse excess peers by closing.
50                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        // One absolute deadline for obtaining the next complete line, even if a
132        // slow sender trickles bytes. Handler time is outside this deadline.
133        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}