Skip to main content

netget/server/fluent_forward/
mod.rs

1//! One-way Fluent Forward 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::{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                // Forward has no negative-ACK grammar: refuse excess peers by closing.
48                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}