netget/server/fluent_forward/
mod.rs1pub 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::{net::SocketAddr, time::Duration};
17use tokio::{
18 io::{AsyncReadExt, AsyncWriteExt},
19 net::{TcpListener, TcpStream},
20};
21
22pub const WRITE_TIMEOUT: Duration = Duration::from_secs(10);
23pub const READ_TIMEOUT: Duration = Duration::from_secs(30);
24pub struct FluentForwardServer;
25impl FluentForwardServer {
26 pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
27 let llm_fallback = ctx
28 .startup_params
29 .as_ref()
30 .map(|p| p.get_optional_bool("llm_fallback"))
31 .transpose()?
32 .flatten()
33 .unwrap_or(codec::DEFAULT_LLM_FALLBACK);
34 let listener = TcpListener::bind(ctx.legacy_listen_addr()).await?;
35 let local = listener.local_addr()?;
36 console_info!(
37 ctx.status_tx,
38 "FluentForward TCP listening on {} (llm_fallback={})",
39 local,
40 llm_fallback
41 );
42 let registrar = ctx.state.clone();
43 let server_id = ctx.server_id;
44 let task = tokio::spawn(async move {
45 let limiter = ConnectionLimiter::new(DEFAULT_MAX_CONNECTIONS);
46 loop {
47 let (socket, peer, permit) = match accept_bounded(
49 &listener,
50 &limiter,
51 b"",
52 "FluentForward",
53 Some(&ctx.status_tx),
54 )
55 .await
56 {
57 Ok(accepted) => accepted,
58 Err(error) => {
59 console_error!(ctx.status_tx, "FluentForward accept failed: {}", error);
60 break;
61 }
62 };
63 let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
64 let now = crate::utils::clock::Instant::now();
65 ctx.state
66 .add_connection_to_server(
67 server_id,
68 ConnectionState {
69 id,
70 remote_addr: peer,
71 local_addr: local,
72 bytes_sent: 0,
73 bytes_received: 0,
74 packets_sent: 0,
75 packets_received: 0,
76 last_activity: now,
77 status: ConnectionStatus::Active,
78 status_changed_at: now,
79 protocol_info: ProtocolConnectionInfo::empty(),
80 },
81 )
82 .await;
83 let child_ctx = ctx.clone();
84 let child = tokio::spawn(async move {
85 let _permit = permit;
86 if let Err(error) = session(socket, peer, id, &child_ctx, llm_fallback).await {
87 console_error!(
88 child_ctx.status_tx,
89 "FluentForward peer {} closed: {}",
90 peer,
91 error
92 );
93 child_ctx
94 .state
95 .record_access_log(
96 AccessLogOwner::Server(server_id.as_u32()),
97 "FluentForward",
98 Some(id.as_u32()),
99 "forward_invalid_stream",
100 json!({"source_addr":peer.to_string()}),
101 vec![json!({"error":error.to_string()})],
102 )
103 .await;
104 }
105 child_ctx
106 .state
107 .remove_connection_from_server(server_id, id)
108 .await;
109 let _ = child_ctx.status_tx.send("__UPDATE_UI__".into());
110 });
111 ctx.state.register_server_task(server_id, child).await;
112 }
113 });
114 registrar.register_server_task(server_id, task).await;
115 Ok(local)
116 }
117}
118async fn session(
119 mut socket: TcpStream,
120 peer: SocketAddr,
121 id: ConnectionId,
122 ctx: &SpawnContext,
123 llm_fallback: bool,
124) -> Result<()> {
125 let mut decoder = codec::Decoder::default();
126 let mut buffer = [0; 8192];
127 let protocol = actions::FluentForwardProtocol::new();
128 loop {
129 let deadline = tokio::time::Instant::now() + READ_TIMEOUT;
130 let batch = loop {
131 if let Some(node) = decoder.next_node()? {
132 if let Some(batch) = codec::parse_batch(node)? {
133 break batch;
134 } else {
135 continue;
136 }
137 }
138 let n = tokio::time::timeout_at(deadline, socket.read(&mut buffer))
139 .await
140 .context("Forward frame read deadline exceeded")??;
141 if n == 0 {
142 return decoder.finish();
143 }
144 ctx.state
145 .update_connection_stats(ctx.server_id, id, Some(n as u64), None, Some(1), None)
146 .await;
147 decoder.feed(&buffer[..n])?;
148 };
149 let event = Event::new(
150 &actions::FORWARD_BATCH_EVENT,
151 json!({"tag":batch.tag,"entries":batch.entries,"mode":batch.mode,"record_count":batch.entries.len(),"ack_requested":batch.chunk.is_some(),"source_addr":peer.to_string()}),
152 );
153 let configured = ctx
154 .state
155 .get_event_handler_config(ctx.server_id)
156 .await
157 .is_some_and(|c| c.find_handler("forward_batch").is_some());
158 if llm_fallback || configured {
159 let result = crate::llm::action_helper::call_llm(
160 &ctx.llm_client,
161 &ctx.state,
162 ctx.server_id,
163 Some(id),
164 &event,
165 &protocol,
166 )
167 .await;
168 let result = match result {
169 Ok(result) => result,
170 Err(error) => {
171 console_error!(
172 ctx.status_tx,
173 "Forward decision=fail_closed_handler_error: {}",
174 error
175 );
176 ctx.state.record_access_log(
177 AccessLogOwner::Server(ctx.server_id.as_u32()), "FluentForward", Some(id.as_u32()),
178 "forward_handler_failed", event.data,
179 vec![json!({"decision":"fail_closed_handler_error","error":error.to_string()})],
180 ).await;
181 return Err(error.context("Forward handler failed"));
182 }
183 };
184 if !result.failures.is_empty() {
185 console_error!(
186 ctx.status_tx,
187 "Forward decision=fail_closed_handler_action_error failed_actions={}",
188 result.failures.len()
189 );
190 ctx.state.record_access_log(
191 AccessLogOwner::Server(ctx.server_id.as_u32()), "FluentForward", Some(id.as_u32()),
192 "forward_handler_failed", event.data,
193 vec![json!({"decision":"fail_closed_handler_action_error","failed_action_count":result.failures.len()})],
194 ).await;
195 anyhow::bail!("Forward handler action failed");
196 }
197 if result
198 .protocol_results
199 .iter()
200 .any(|r| r.closes_connection())
201 {
202 return Ok(());
203 }
204 for message in result.messages {
205 console_info!(ctx.status_tx, "{}", message);
206 }
207 } else {
208 ctx.state
209 .record_access_log(
210 AccessLogOwner::Server(ctx.server_id.as_u32()),
211 "FluentForward",
212 Some(id.as_u32()),
213 "forward_batch",
214 event.data,
215 vec![json!({"type":"accept_forward_batch"})],
216 )
217 .await;
218 }
219 if let Some(chunk) = batch.chunk {
220 let bytes = codec::encode_ack(&chunk)?;
221 tokio::time::timeout(WRITE_TIMEOUT, socket.write_all(&bytes))
222 .await
223 .context("Forward ACK write deadline exceeded")??;
224 ctx.state
225 .update_connection_stats(
226 ctx.server_id,
227 id,
228 None,
229 Some(bytes.len() as u64),
230 None,
231 Some(1),
232 )
233 .await;
234 }
235 let _ = ctx.status_tx.send("__UPDATE_UI__".into());
236 }
237}