nemo_relay/observability/
atof.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Agent Trajectory Observability Format (ATOF) JSONL exporter support for NeMo
5//! Flow.
6//!
7//! The [`AtofExporter`] registers as an event subscriber and writes each
8//! canonical NeMo Relay Agent Trajectory Observability Format (ATOF) event as
9//! one JSON object per JSONL line.
10
11use std::collections::HashMap;
12use std::fs::{File, OpenOptions};
13use std::io::{BufWriter, Write};
14use std::path::{Path, PathBuf};
15use std::sync::{Arc, Mutex, mpsc as std_mpsc};
16use std::time::Duration;
17
18use chrono::Utc;
19#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
20use futures_util::{SinkExt, stream};
21use serde::{Deserialize, Serialize};
22use serde_json::Value as Json;
23#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
24use tokio_tungstenite::tungstenite::client::IntoClientRequest;
25
26use crate::api::event::Event;
27use crate::api::runtime::EventSubscriberFn;
28use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber};
29use crate::error::FlowError;
30
31/// Result type for the ATOF JSONL exporter.
32pub type Result<T> = std::result::Result<T, AtofExporterError>;
33
34/// Errors produced while configuring or operating the ATOF JSONL exporter.
35#[derive(Debug, thiserror::Error)]
36pub enum AtofExporterError {
37    /// Failed to resolve the current working directory for default config.
38    #[error("failed to resolve current working directory: {0}")]
39    CurrentDirectory(std::io::Error),
40    /// Failed to open the output file.
41    #[error("failed to open ATOF output file {path:?}: {source}")]
42    OpenFile {
43        /// Output path that failed to open.
44        path: PathBuf,
45        /// Underlying I/O error.
46        source: std::io::Error,
47    },
48    /// Failed while flushing the output file.
49    #[error("failed to flush ATOF output file {path:?}: {source}")]
50    Flush {
51        /// Output path that failed to flush.
52        path: PathBuf,
53        /// Underlying I/O error.
54        source: std::io::Error,
55    },
56    /// The exporter recorded an earlier write or serialization error.
57    #[error("previous ATOF export failed for {path:?}: {message}")]
58    StoredFailure {
59        /// Output path associated with the failure.
60        path: PathBuf,
61        /// Stored failure message.
62        message: String,
63    },
64    /// A streaming endpoint configuration is invalid.
65    #[error("invalid ATOF streaming endpoint: {0}")]
66    InvalidEndpoint(String),
67    /// The internal exporter state lock was poisoned.
68    #[error("the ATOF exporter state lock was poisoned")]
69    LockPoisoned,
70    /// Runtime subscriber registration failed.
71    #[error(transparent)]
72    Runtime(#[from] FlowError),
73}
74
75/// File write behavior for [`AtofExporter`].
76#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
77#[serde(rename_all = "snake_case")]
78pub enum AtofExporterMode {
79    /// Append events to an existing file or create it if missing.
80    #[default]
81    Append,
82    /// Truncate an existing file when the exporter is created.
83    Overwrite,
84}
85
86impl AtofExporterMode {
87    /// Parse a string mode used by language bindings.
88    pub fn parse(value: &str) -> Option<Self> {
89        match value {
90            "append" => Some(Self::Append),
91            "overwrite" => Some(Self::Overwrite),
92            _ => None,
93        }
94    }
95
96    /// Return the stable string representation used by language bindings.
97    pub fn as_str(self) -> &'static str {
98        match self {
99            Self::Append => "append",
100            Self::Overwrite => "overwrite",
101        }
102    }
103}
104
105/// Streaming transport used by an ATOF endpoint.
106#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
107#[serde(rename_all = "snake_case")]
108pub enum AtofEndpointTransport {
109    /// POST each event as one JSONL record.
110    #[default]
111    HttpPost,
112    /// Send each event as one WebSocket JSON text message.
113    Websocket,
114    /// Stream events over one long-lived HTTP NDJSON upload.
115    Ndjson,
116}
117
118impl AtofEndpointTransport {
119    /// Parse a string transport used by configuration and bindings.
120    pub fn parse(value: &str) -> Option<Self> {
121        match value {
122            "http_post" => Some(Self::HttpPost),
123            "websocket" => Some(Self::Websocket),
124            "ndjson" => Some(Self::Ndjson),
125            _ => None,
126        }
127    }
128
129    /// Return the stable string representation used by configuration and bindings.
130    pub fn as_str(self) -> &'static str {
131        match self {
132            Self::HttpPost => "http_post",
133            Self::Websocket => "websocket",
134            Self::Ndjson => "ndjson",
135        }
136    }
137}
138
139/// Streaming destination for raw ATOF events.
140#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
141pub struct AtofEndpointConfig {
142    /// Endpoint URL.
143    pub url: String,
144    /// Endpoint transport.
145    #[serde(default)]
146    pub transport: AtofEndpointTransport,
147    /// Headers applied to endpoint requests or handshakes.
148    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
149    pub headers: HashMap<String, String>,
150    /// Per-endpoint timeout in milliseconds.
151    #[serde(default = "default_endpoint_timeout_millis")]
152    pub timeout_millis: u64,
153}
154
155impl AtofEndpointConfig {
156    /// Create a streaming endpoint with defaults.
157    pub fn new(url: impl Into<String>, transport: AtofEndpointTransport) -> Self {
158        Self {
159            url: url.into(),
160            transport,
161            headers: HashMap::new(),
162            timeout_millis: default_endpoint_timeout_millis(),
163        }
164    }
165
166    /// Add a header to this endpoint config.
167    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
168        self.headers.insert(key.into(), value.into());
169        self
170    }
171
172    /// Override the endpoint timeout.
173    pub fn with_timeout_millis(mut self, timeout_millis: u64) -> Self {
174        self.timeout_millis = timeout_millis;
175        self
176    }
177}
178
179/// Configuration for [`AtofExporter`].
180#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
181pub struct AtofExporterConfig {
182    /// Directory that contains the JSONL output file.
183    #[serde(default = "default_output_directory")]
184    pub output_directory: PathBuf,
185    /// Append or overwrite behavior used when opening the file.
186    #[serde(default)]
187    pub mode: AtofExporterMode,
188    /// Output filename.
189    #[serde(default = "default_filename")]
190    pub filename: String,
191    /// Optional streaming endpoints that receive every raw ATOF event.
192    #[serde(default, skip_serializing_if = "Vec::is_empty")]
193    pub endpoints: Vec<AtofEndpointConfig>,
194}
195
196impl Default for AtofExporterConfig {
197    fn default() -> Self {
198        Self {
199            output_directory: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
200            mode: AtofExporterMode::Append,
201            filename: default_filename(),
202            endpoints: Vec::new(),
203        }
204    }
205}
206
207impl AtofExporterConfig {
208    /// Create a config with defaults.
209    pub fn new() -> Self {
210        Self::default()
211    }
212
213    /// Override the output directory.
214    pub fn with_output_directory(mut self, output_directory: impl Into<PathBuf>) -> Self {
215        self.output_directory = output_directory.into();
216        self
217    }
218
219    /// Override the output mode.
220    pub fn with_mode(mut self, mode: AtofExporterMode) -> Self {
221        self.mode = mode;
222        self
223    }
224
225    /// Override the output filename.
226    pub fn with_filename(mut self, filename: impl Into<String>) -> Self {
227        self.filename = filename.into();
228        self
229    }
230
231    /// Override streaming endpoints.
232    pub fn with_endpoints(mut self, endpoints: Vec<AtofEndpointConfig>) -> Self {
233        self.endpoints = endpoints;
234        self
235    }
236
237    /// Add one streaming endpoint.
238    pub fn with_endpoint(mut self, endpoint: AtofEndpointConfig) -> Self {
239        self.endpoints.push(endpoint);
240        self
241    }
242
243    /// Return the full output path for this config.
244    pub fn path(&self) -> PathBuf {
245        self.output_directory.join(&self.filename)
246    }
247}
248
249struct AtofExporterState {
250    writer: BufWriter<File>,
251    last_error: Option<String>,
252    endpoints: Vec<AtofEndpointWorker>,
253    closed: bool,
254}
255
256/// Filesystem-backed Agent Trajectory Observability Format (ATOF) JSONL event exporter.
257pub struct AtofExporter {
258    path: PathBuf,
259    state: Arc<Mutex<AtofExporterState>>,
260}
261
262impl AtofExporter {
263    /// Create a new exporter from config and open its output file.
264    pub fn new(config: AtofExporterConfig) -> Result<Self> {
265        let path = config.path();
266        let file = open_file(&path, config.mode)?;
267        let endpoints = start_endpoint_workers(&config.endpoints)?;
268        Ok(Self {
269            path,
270            state: Arc::new(Mutex::new(AtofExporterState {
271                writer: BufWriter::new(file),
272                last_error: None,
273                endpoints,
274                closed: false,
275            })),
276        })
277    }
278
279    /// Return the output JSONL path.
280    pub fn path(&self) -> &Path {
281        self.path.as_path()
282    }
283
284    /// Return an event subscriber that writes one JSONL record per observed event.
285    pub fn subscriber(&self) -> EventSubscriberFn {
286        let state = Arc::clone(&self.state);
287        Arc::new(move |event: &Event| {
288            let Ok(mut state) = state.lock() else {
289                return;
290            };
291            if state.closed || state.last_error.is_some() {
292                return;
293            }
294            let Ok(value) = event.try_to_json_value() else {
295                state.last_error = Some("failed to serialize ATOF event".to_string());
296                return;
297            };
298            if let Err(error) = write_json_value(&mut state.writer, &value) {
299                state.last_error = Some(error);
300                return;
301            }
302            let Ok(raw_json) = serde_json::to_string(&value) else {
303                state.last_error = Some("failed to serialize ATOF event".to_string());
304                return;
305            };
306            for endpoint in &state.endpoints {
307                endpoint.enqueue(raw_json.clone());
308            }
309        })
310    }
311
312    /// Register this exporter globally under the given subscriber name.
313    pub fn register(&self, name: &str) -> Result<()> {
314        register_subscriber(name, self.subscriber()).map_err(Into::into)
315    }
316
317    /// Deregister a global subscriber by name.
318    pub fn deregister(&self, name: &str) -> Result<bool> {
319        deregister_subscriber(name).map_err(Into::into)
320    }
321
322    /// Flush the underlying file and drain queued endpoint events.
323    pub fn force_flush(&self) -> Result<()> {
324        flush_subscribers()?;
325        let mut state = self
326            .state
327            .lock()
328            .map_err(|_| AtofExporterError::LockPoisoned)?;
329        if state.closed {
330            return stored_failure_result(&self.path, &state);
331        }
332        state
333            .writer
334            .flush()
335            .map_err(|source| AtofExporterError::Flush {
336                path: self.path.clone(),
337                source,
338            })?;
339        for endpoint in &state.endpoints {
340            endpoint.flush();
341        }
342        stored_failure_result(&self.path, &state)
343    }
344
345    /// Shut down the exporter by flushing buffered data and closing endpoints.
346    pub fn shutdown(&self) -> Result<()> {
347        flush_subscribers()?;
348        let mut state = self
349            .state
350            .lock()
351            .map_err(|_| AtofExporterError::LockPoisoned)?;
352        if state.closed {
353            return stored_failure_result(&self.path, &state);
354        }
355        state.closed = true;
356        let flush_result = state
357            .writer
358            .flush()
359            .map_err(|source| AtofExporterError::Flush {
360                path: self.path.clone(),
361                source,
362            });
363        for endpoint in &state.endpoints {
364            endpoint.close();
365        }
366        flush_result?;
367        stored_failure_result(&self.path, &state)
368    }
369}
370
371fn default_filename() -> String {
372    format!(
373        "nemo-relay-events-{}.jsonl",
374        Utc::now().format("%Y-%m-%d-%H.%M.%S")
375    )
376}
377
378fn default_output_directory() -> PathBuf {
379    std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
380}
381
382fn default_endpoint_timeout_millis() -> u64 {
383    3_000
384}
385
386fn open_file(path: &Path, mode: AtofExporterMode) -> Result<File> {
387    let mut options = OpenOptions::new();
388    options.create(true);
389    match mode {
390        AtofExporterMode::Append => {
391            options.append(true);
392        }
393        AtofExporterMode::Overwrite => {
394            options.write(true).truncate(true);
395        }
396    }
397    options
398        .open(path)
399        .map_err(|source| AtofExporterError::OpenFile {
400            path: path.to_path_buf(),
401            source,
402        })
403}
404
405fn write_json_value(writer: &mut BufWriter<File>, value: &Json) -> std::result::Result<(), String> {
406    serde_json::to_writer(&mut *writer, value).map_err(|error| error.to_string())?;
407    writer.write_all(b"\n").map_err(|error| error.to_string())?;
408    writer.flush().map_err(|error| error.to_string())
409}
410
411fn stored_failure_result(path: &Path, state: &AtofExporterState) -> Result<()> {
412    if let Some(message) = &state.last_error {
413        return Err(AtofExporterError::StoredFailure {
414            path: path.to_path_buf(),
415            message: message.clone(),
416        });
417    }
418    Ok(())
419}
420
421#[cfg_attr(
422    any(not(feature = "atof-streaming"), target_arch = "wasm32"),
423    allow(dead_code)
424)]
425enum EndpointMessage {
426    Event(String),
427    Flush(std_mpsc::Sender<()>),
428    Close(std_mpsc::Sender<()>),
429}
430
431#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
432enum NdjsonBodyMessage {
433    Event(Vec<u8>),
434    Flush(std_mpsc::Sender<()>),
435}
436
437#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
438impl NdjsonBodyMessage {
439    fn acknowledge_if_flush(self) {
440        if let Self::Flush(done) = self {
441            let _ = done.send(());
442        }
443    }
444}
445
446struct AtofEndpointWorker {
447    sender: tokio::sync::mpsc::UnboundedSender<EndpointMessage>,
448    timeout: Duration,
449}
450
451impl AtofEndpointWorker {
452    fn enqueue(&self, raw_json: String) {
453        let _ = self.sender.send(EndpointMessage::Event(raw_json));
454    }
455
456    fn flush(&self) {
457        let (tx, rx) = std_mpsc::channel();
458        if self.sender.send(EndpointMessage::Flush(tx)).is_ok()
459            && rx.recv_timeout(self.timeout).is_err()
460        {
461            eprintln!("nemo_relay: timed out flushing ATOF endpoint");
462        }
463    }
464
465    fn close(&self) {
466        let (tx, rx) = std_mpsc::channel();
467        if self.sender.send(EndpointMessage::Close(tx)).is_ok()
468            && rx.recv_timeout(self.timeout).is_err()
469        {
470            eprintln!("nemo_relay: timed out closing ATOF endpoint");
471        }
472    }
473}
474
475#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
476fn start_endpoint_workers(configs: &[AtofEndpointConfig]) -> Result<Vec<AtofEndpointWorker>> {
477    let mut workers = Vec::with_capacity(configs.len());
478    for (index, config) in configs.iter().enumerate() {
479        match start_endpoint_worker(index, config.clone()) {
480            Ok(worker) => workers.push(worker),
481            Err(error) => {
482                eprintln!("nemo_relay: invalid ATOF endpoint[{index}]: {error}");
483                return Err(AtofExporterError::InvalidEndpoint(format!(
484                    "endpoints[{index}]: {error}"
485                )));
486            }
487        }
488    }
489    Ok(workers)
490}
491
492#[cfg(any(not(feature = "atof-streaming"), target_arch = "wasm32"))]
493fn start_endpoint_workers(configs: &[AtofEndpointConfig]) -> Result<Vec<AtofEndpointWorker>> {
494    if configs.is_empty() {
495        return Ok(Vec::new());
496    }
497    let message = "ATOF streaming endpoints are not supported in this build".to_string();
498    eprintln!("nemo_relay: {message}");
499    Err(AtofExporterError::InvalidEndpoint(message))
500}
501
502#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
503fn start_endpoint_worker(index: usize, config: AtofEndpointConfig) -> Result<AtofEndpointWorker> {
504    validate_endpoint_config(&config)?;
505    let timeout = Duration::from_millis(config.timeout_millis);
506    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
507    std::thread::Builder::new()
508        .name(format!("nemo-relay-atof-endpoint-{index}"))
509        .spawn(move || run_endpoint_worker(index, config, rx))
510        .map_err(|error| AtofExporterError::InvalidEndpoint(error.to_string()))?;
511    Ok(AtofEndpointWorker {
512        sender: tx,
513        timeout,
514    })
515}
516
517#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
518fn validate_endpoint_config(config: &AtofEndpointConfig) -> Result<()> {
519    if config.url.trim().is_empty() {
520        return Err(AtofExporterError::InvalidEndpoint(
521            "endpoint url must be non-empty".to_string(),
522        ));
523    }
524    if config.timeout_millis == 0 {
525        return Err(AtofExporterError::InvalidEndpoint(
526            "endpoint timeout_millis must be greater than 0".to_string(),
527        ));
528    }
529    let url = reqwest::Url::parse(&config.url)
530        .map_err(|error| AtofExporterError::InvalidEndpoint(error.to_string()))?;
531    let valid_scheme = match config.transport {
532        AtofEndpointTransport::HttpPost | AtofEndpointTransport::Ndjson => {
533            matches!(url.scheme(), "http" | "https")
534        }
535        AtofEndpointTransport::Websocket => matches!(url.scheme(), "ws" | "wss"),
536    };
537    if !valid_scheme {
538        return Err(AtofExporterError::InvalidEndpoint(format!(
539            "endpoint {} transport does not support URL scheme {:?}",
540            config.transport.as_str(),
541            url.scheme()
542        )));
543    }
544    build_header_map(&config.headers)?;
545    Ok(())
546}
547
548#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
549fn build_header_map(headers: &HashMap<String, String>) -> Result<reqwest::header::HeaderMap> {
550    let mut out = reqwest::header::HeaderMap::new();
551    for (key, value) in headers {
552        let name = reqwest::header::HeaderName::from_bytes(key.as_bytes())
553            .map_err(|error| AtofExporterError::InvalidEndpoint(error.to_string()))?;
554        let value = reqwest::header::HeaderValue::from_str(value)
555            .map_err(|error| AtofExporterError::InvalidEndpoint(error.to_string()))?;
556        out.insert(name, value);
557    }
558    Ok(out)
559}
560
561#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
562fn run_endpoint_worker(
563    index: usize,
564    config: AtofEndpointConfig,
565    rx: tokio::sync::mpsc::UnboundedReceiver<EndpointMessage>,
566) {
567    let runtime = match tokio::runtime::Builder::new_current_thread()
568        .enable_all()
569        .build()
570    {
571        Ok(runtime) => runtime,
572        Err(error) => {
573            eprintln!("nemo_relay: ATOF endpoint[{index}] runtime failed: {error}");
574            return;
575        }
576    };
577    runtime.block_on(async move {
578        match config.transport {
579            AtofEndpointTransport::HttpPost => run_http_post_endpoint(index, config, rx).await,
580            AtofEndpointTransport::Websocket => run_websocket_endpoint(index, config, rx).await,
581            AtofEndpointTransport::Ndjson => run_ndjson_endpoint(index, config, rx).await,
582        }
583    });
584}
585
586#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
587async fn run_http_post_endpoint(
588    index: usize,
589    config: AtofEndpointConfig,
590    mut rx: tokio::sync::mpsc::UnboundedReceiver<EndpointMessage>,
591) {
592    let client = match reqwest::Client::builder()
593        .timeout(Duration::from_millis(config.timeout_millis))
594        .default_headers(match build_header_map(&config.headers) {
595            Ok(headers) => headers,
596            Err(error) => {
597                eprintln!("nemo_relay: ATOF endpoint[{index}] disabled: {error}");
598                drain_closed(rx).await;
599                return;
600            }
601        })
602        .build()
603    {
604        Ok(client) => client,
605        Err(error) => {
606            eprintln!("nemo_relay: ATOF endpoint[{index}] client build failed: {error}");
607            drain_closed(rx).await;
608            return;
609        }
610    };
611    while let Some(message) = rx.recv().await {
612        match message {
613            EndpointMessage::Event(raw_json) => {
614                let body = format!("{raw_json}\n");
615                let result = client
616                    .post(&config.url)
617                    .header(reqwest::header::CONTENT_TYPE, "application/x-ndjson")
618                    .body(body)
619                    .send()
620                    .await;
621                match result {
622                    Ok(response) if response.status().is_success() => {}
623                    Ok(response) => eprintln!(
624                        "nemo_relay: ATOF endpoint[{index}] HTTP status {}",
625                        response.status()
626                    ),
627                    Err(error) => {
628                        eprintln!("nemo_relay: ATOF endpoint[{index}] send failed: {error}")
629                    }
630                }
631            }
632            EndpointMessage::Flush(done) => {
633                let _ = done.send(());
634            }
635            EndpointMessage::Close(done) => {
636                let _ = done.send(());
637                return;
638            }
639        }
640    }
641}
642
643#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
644async fn run_websocket_endpoint(
645    index: usize,
646    config: AtofEndpointConfig,
647    mut rx: tokio::sync::mpsc::UnboundedReceiver<EndpointMessage>,
648) {
649    let mut pending = std::collections::VecDeque::new();
650    let mut socket = match connect_websocket(&config).await {
651        Ok(socket) => Some(socket),
652        Err(error) => {
653            eprintln!("nemo_relay: ATOF endpoint[{index}] websocket startup failed: {error}");
654            None
655        }
656    };
657    while let Some(message) = rx.recv().await {
658        match message {
659            EndpointMessage::Event(raw_json) => {
660                pending.push_back(raw_json);
661                let _ = drain_websocket_pending(index, &config, &mut socket, &mut pending).await;
662            }
663            EndpointMessage::Flush(done) => {
664                let _ = drain_websocket_pending(index, &config, &mut socket, &mut pending).await;
665                let _ = done.send(());
666            }
667            EndpointMessage::Close(done) => {
668                let _ = drain_websocket_pending(index, &config, &mut socket, &mut pending).await;
669                if let Some(mut ws) = socket.take() {
670                    let _ = ws.close(None).await;
671                }
672                let _ = done.send(());
673                return;
674            }
675        }
676    }
677}
678
679#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
680type AtofWebSocket =
681    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
682
683#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
684async fn drain_websocket_pending(
685    index: usize,
686    config: &AtofEndpointConfig,
687    socket: &mut Option<AtofWebSocket>,
688    pending: &mut std::collections::VecDeque<String>,
689) -> bool {
690    let timeout = Duration::from_millis(config.timeout_millis);
691    match tokio::time::timeout(
692        timeout,
693        drain_websocket_pending_inner(index, config, socket, pending),
694    )
695    .await
696    {
697        Ok(drained) => drained,
698        Err(_) => {
699            eprintln!("nemo_relay: ATOF endpoint[{index}] websocket drain timed out");
700            false
701        }
702    }
703}
704
705#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
706async fn drain_websocket_pending_inner(
707    index: usize,
708    config: &AtofEndpointConfig,
709    socket: &mut Option<AtofWebSocket>,
710    pending: &mut std::collections::VecDeque<String>,
711) -> bool {
712    while let Some(raw_json) = pending.front().cloned() {
713        if socket.is_none() {
714            match connect_websocket(config).await {
715                Ok(ws) => *socket = Some(ws),
716                Err(error) => {
717                    eprintln!(
718                        "nemo_relay: ATOF endpoint[{index}] websocket reconnect failed: {error}"
719                    );
720                    tokio::time::sleep(Duration::from_millis(50)).await;
721                    continue;
722                }
723            }
724        }
725
726        let Some(ws) = socket.as_mut() else {
727            continue;
728        };
729        match ws
730            .send(tokio_tungstenite::tungstenite::Message::Text(
731                raw_json.into(),
732            ))
733            .await
734        {
735            Ok(()) => {
736                pending.pop_front();
737            }
738            Err(error) => {
739                eprintln!("nemo_relay: ATOF endpoint[{index}] websocket send failed: {error}");
740                *socket = None;
741                tokio::time::sleep(Duration::from_millis(50)).await;
742            }
743        }
744    }
745    true
746}
747
748#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
749async fn connect_websocket(
750    config: &AtofEndpointConfig,
751) -> std::result::Result<AtofWebSocket, String> {
752    let mut request = config
753        .url
754        .as_str()
755        .into_client_request()
756        .map_err(|error| error.to_string())?;
757    for (key, value) in &config.headers {
758        let name =
759            tokio_tungstenite::tungstenite::http::header::HeaderName::from_bytes(key.as_bytes())
760                .map_err(|error| error.to_string())?;
761        let value = tokio_tungstenite::tungstenite::http::header::HeaderValue::from_str(value)
762            .map_err(|error| error.to_string())?;
763        request.headers_mut().insert(name, value);
764    }
765    tokio::time::timeout(
766        Duration::from_millis(config.timeout_millis),
767        tokio_tungstenite::connect_async(request),
768    )
769    .await
770    .map_err(|error| error.to_string())?
771    .map(|(socket, _)| socket)
772    .map_err(|error| error.to_string())
773}
774
775#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
776async fn run_ndjson_endpoint(
777    index: usize,
778    config: AtofEndpointConfig,
779    mut rx: tokio::sync::mpsc::UnboundedReceiver<EndpointMessage>,
780) {
781    let client = match build_ndjson_client(&config) {
782        Ok(client) => client,
783        Err(error) => {
784            eprintln!("nemo_relay: ATOF endpoint[{index}] {error}");
785            drain_closed(rx).await;
786            return;
787        }
788    };
789
790    let (body_tx, body) = ndjson_body_channel();
791    let request = tokio::spawn(async move {
792        client
793            .post(config.url)
794            .header(reqwest::header::CONTENT_TYPE, "application/x-ndjson")
795            .body(body)
796            .send()
797            .await
798    });
799    let close_timeout = Duration::from_millis(config.timeout_millis);
800
801    while let Some(message) = rx.recv().await {
802        match message {
803            EndpointMessage::Event(raw_json) => send_ndjson_event(index, &body_tx, raw_json),
804            EndpointMessage::Flush(done) => send_ndjson_flush(index, &body_tx, done),
805            EndpointMessage::Close(done) => {
806                drop(body_tx);
807                finish_ndjson_upload(index, request, close_timeout, done).await;
808                return;
809            }
810        }
811    }
812}
813
814#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
815fn build_ndjson_client(
816    config: &AtofEndpointConfig,
817) -> std::result::Result<reqwest::Client, String> {
818    let headers =
819        build_header_map(&config.headers).map_err(|error| format!("disabled: {error}"))?;
820    reqwest::Client::builder()
821        .connect_timeout(Duration::from_millis(config.timeout_millis))
822        .default_headers(headers)
823        .build()
824        .map_err(|error| format!("client build failed: {error}"))
825}
826
827#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
828fn ndjson_body_channel() -> (
829    tokio::sync::mpsc::UnboundedSender<NdjsonBodyMessage>,
830    reqwest::Body,
831) {
832    let (body_tx, body_rx) = tokio::sync::mpsc::unbounded_channel::<NdjsonBodyMessage>();
833    let body_stream = stream::unfold(body_rx, |mut body_rx| async {
834        loop {
835            match body_rx.recv().await? {
836                NdjsonBodyMessage::Event(bytes) => {
837                    return Some((Ok::<_, std::io::Error>(bytes), body_rx));
838                }
839                NdjsonBodyMessage::Flush(done) => {
840                    let _ = done.send(());
841                }
842            }
843        }
844    });
845    (body_tx, reqwest::Body::wrap_stream(body_stream))
846}
847
848#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
849fn send_ndjson_event(
850    index: usize,
851    body_tx: &tokio::sync::mpsc::UnboundedSender<NdjsonBodyMessage>,
852    raw_json: String,
853) {
854    if let Err(error) = body_tx.send(NdjsonBodyMessage::Event(
855        format!("{raw_json}\n").into_bytes(),
856    )) {
857        eprintln!("nemo_relay: ATOF endpoint[{index}] NDJSON send failed: {error}");
858    }
859}
860
861#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
862fn send_ndjson_flush(
863    index: usize,
864    body_tx: &tokio::sync::mpsc::UnboundedSender<NdjsonBodyMessage>,
865    done: std_mpsc::Sender<()>,
866) {
867    if let Err(error) = body_tx.send(NdjsonBodyMessage::Flush(done)) {
868        eprintln!("nemo_relay: ATOF endpoint[{index}] NDJSON flush failed: {error}");
869        error.0.acknowledge_if_flush();
870    }
871}
872
873#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
874async fn finish_ndjson_upload(
875    index: usize,
876    request: tokio::task::JoinHandle<reqwest::Result<reqwest::Response>>,
877    close_timeout: Duration,
878    done: std_mpsc::Sender<()>,
879) {
880    match tokio::time::timeout(close_timeout, request).await {
881        Ok(Ok(Ok(response))) if response.status().is_success() => {}
882        Ok(Ok(Ok(response))) => eprintln!(
883            "nemo_relay: ATOF endpoint[{index}] NDJSON HTTP status {}",
884            response.status()
885        ),
886        Ok(Ok(Err(error))) => {
887            eprintln!("nemo_relay: ATOF endpoint[{index}] NDJSON upload failed: {error}")
888        }
889        Ok(Err(error)) => {
890            eprintln!("nemo_relay: ATOF endpoint[{index}] NDJSON task failed: {error}")
891        }
892        Err(_) => eprintln!("nemo_relay: ATOF endpoint[{index}] NDJSON close timed out"),
893    }
894    let _ = done.send(());
895}
896
897#[cfg(all(feature = "atof-streaming", not(target_arch = "wasm32")))]
898async fn drain_closed(mut rx: tokio::sync::mpsc::UnboundedReceiver<EndpointMessage>) {
899    while let Some(message) = rx.recv().await {
900        match message {
901            EndpointMessage::Flush(done) => {
902                let _ = done.send(());
903            }
904            EndpointMessage::Close(done) => {
905                let _ = done.send(());
906                return;
907            }
908            EndpointMessage::Event(_) => {}
909        }
910    }
911}
912
913// ---------------------------------------------------------------------------
914// Tests
915// ---------------------------------------------------------------------------
916
917#[cfg(test)]
918#[path = "../../tests/unit/observability/atof_tests.rs"]
919mod tests;
nemo_relay/observability/atof.rs

nemo_relay/observability/
atof.rs
