Skip to main content

hyperactor_mesh/
logging.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use std::collections::HashMap;
10use std::collections::VecDeque;
11use std::fmt;
12use std::path::Path;
13use std::path::PathBuf;
14use std::sync::Arc;
15use std::time::Duration;
16use std::time::SystemTime;
17
18use anyhow::Result;
19use async_trait::async_trait;
20use chrono::DateTime;
21use chrono::Local;
22use hostname;
23use hyperactor::Actor;
24use hyperactor::ActorEnvironment;
25use hyperactor::ActorRef;
26use hyperactor::Context;
27use hyperactor::Endpoint as _;
28use hyperactor::HandleClient;
29use hyperactor::Handler;
30use hyperactor::Instance;
31use hyperactor::OncePortRef;
32use hyperactor::ProcAddr;
33use hyperactor::RefClient;
34use hyperactor::channel;
35use hyperactor::channel::ChannelAddr;
36use hyperactor::channel::ChannelRx;
37use hyperactor::channel::ChannelTransport;
38use hyperactor::channel::ChannelTx;
39use hyperactor::channel::Rx;
40use hyperactor::channel::Tx;
41use hyperactor::channel::TxStatus;
42use hyperactor_config::CONFIG;
43use hyperactor_config::ConfigAttr;
44use hyperactor_config::attrs::declare_attrs;
45use hyperactor_telemetry::env;
46use hyperactor_telemetry::log_file_path;
47use serde::Deserialize;
48use serde::Serialize;
49use tokio::io;
50use tokio::io::AsyncRead;
51use tokio::io::AsyncReadExt;
52use tokio::io::AsyncWriteExt;
53use tokio::sync::Mutex;
54use tokio::sync::Notify;
55use tokio::sync::RwLock;
56use tokio::sync::watch::Receiver;
57use tokio::task::JoinHandle;
58use tracing::Level;
59use typeuri::Named;
60
61use crate::bootstrap::BOOTSTRAP_LOG_CHANNEL;
62use crate::shortuuid::ShortUuid;
63
64mod line_prefixing_writer;
65
66pub(crate) const DEFAULT_AGGREGATE_WINDOW_SEC: u64 = 5;
67const MAX_LINE_SIZE: usize = 4 * 1024;
68
69declare_attrs! {
70    /// Maximum number of lines to batch before flushing to client
71    /// This means that stdout/err reader will be paused after reading `HYPERACTOR_READ_LOG_BUFFER` lines.
72    /// After pause lines will be flushed and reading will resume.
73    @meta(CONFIG = ConfigAttr::new(
74        Some("HYPERACTOR_READ_LOG_BUFFER".to_string()),
75        Some("read_log_buffer".to_string()),
76    ))
77    pub attr READ_LOG_BUFFER: usize = 100;
78
79    /// If enabled, local logs are also written to a file and aggregated
80    @meta(CONFIG = ConfigAttr::new(
81        Some("HYPERACTOR_FORCE_FILE_LOG".to_string()),
82        Some("force_file_log".to_string()),
83    ))
84    pub attr FORCE_FILE_LOG: bool = false;
85
86    /// Prefixes logs with rank
87    @meta(CONFIG = ConfigAttr::new(
88        Some("HYPERACTOR_PREFIX_WITH_RANK".to_string()),
89        Some("prefix_with_rank".to_string()),
90    ))
91    pub attr PREFIX_WITH_RANK: bool = true;
92}
93
94/// Calculate the Levenshtein distance between two strings
95fn levenshtein_distance(left: &str, right: &str) -> usize {
96    let left_chars: Vec<char> = left.chars().collect();
97    let right_chars: Vec<char> = right.chars().collect();
98
99    let left_len = left_chars.len();
100    let right_len = right_chars.len();
101
102    // Handle edge cases
103    if left_len == 0 {
104        return right_len;
105    }
106    if right_len == 0 {
107        return left_len;
108    }
109
110    // Create a matrix of size (len_s1+1) x (len_s2+1)
111    let mut matrix = vec![vec![0; right_len + 1]; left_len + 1];
112
113    // Initialize the first row and column
114    for (i, row) in matrix.iter_mut().enumerate().take(left_len + 1) {
115        row[0] = i;
116    }
117    for (j, cell) in matrix[0].iter_mut().enumerate().take(right_len + 1) {
118        *cell = j;
119    }
120
121    // Fill the matrix
122    for i in 1..=left_len {
123        for j in 1..=right_len {
124            let cost = if left_chars[i - 1] == right_chars[j - 1] {
125                0
126            } else {
127                1
128            };
129
130            matrix[i][j] = std::cmp::min(
131                std::cmp::min(
132                    matrix[i - 1][j] + 1, // deletion
133                    matrix[i][j - 1] + 1, // insertion
134                ),
135                matrix[i - 1][j - 1] + cost, // substitution
136            );
137        }
138    }
139
140    // Return the bottom-right cell
141    matrix[left_len][right_len]
142}
143
144/// Calculate the normalized edit distance between two strings (0.0 to 1.0)
145fn normalized_edit_distance(left: &str, right: &str) -> f64 {
146    let distance = levenshtein_distance(left, right) as f64;
147    let max_len = std::cmp::max(left.len(), right.len()) as f64;
148
149    if max_len == 0.0 {
150        0.0 // Both strings are empty, so they're identical
151    } else {
152        distance / max_len
153    }
154}
155
156#[derive(Debug, Clone)]
157/// LogLine represents a single log line with its content and count
158struct LogLine {
159    content: String,
160    pub count: u64,
161}
162
163impl LogLine {
164    fn new(content: String) -> Self {
165        Self { content, count: 1 }
166    }
167}
168
169impl fmt::Display for LogLine {
170    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
171        write!(
172            f,
173            "\x1b[33m[{} similar log lines]\x1b[0m {}",
174            self.count, self.content
175        )
176    }
177}
178
179#[derive(Debug, Clone)]
180/// Aggregator is a struct that holds a list of LogLines and a start time.
181/// It can aggregate new log lines to existing ones if they are "similar" based on edit distance.
182struct Aggregator {
183    lines: Vec<LogLine>,
184    start_time: SystemTime,
185    similarity_threshold: f64, // Threshold for considering two strings similar (0.0 to 1.0)
186}
187
188impl Aggregator {
189    fn new() -> Self {
190        // Default threshold: strings with normalized edit distance < 0.15 are considered similar
191        Self::new_with_threshold(0.15)
192    }
193
194    fn new_with_threshold(threshold: f64) -> Self {
195        Aggregator {
196            lines: vec![],
197            start_time: std::time::SystemTime::now(),
198            similarity_threshold: threshold,
199        }
200    }
201
202    fn reset(&mut self) {
203        self.lines.clear();
204        self.start_time = std::time::SystemTime::now();
205    }
206
207    fn add_line(&mut self, line: &str) -> anyhow::Result<()> {
208        // Find the most similar existing line
209        let mut best_match_idx = None;
210        let mut best_similarity = f64::MAX;
211
212        for (idx, existing_line) in self.lines.iter().enumerate() {
213            let distance = normalized_edit_distance(&existing_line.content, line);
214
215            // If this line is more similar than our current best match
216            if distance < best_similarity && distance < self.similarity_threshold {
217                best_match_idx = Some(idx);
218                best_similarity = distance;
219            }
220        }
221
222        // If we found a similar enough line, increment its count
223        if let Some(idx) = best_match_idx {
224            self.lines[idx].count += 1;
225        } else {
226            // Otherwise, add a new line
227            self.lines.push(LogLine::new(line.to_string()));
228        }
229
230        Ok(())
231    }
232
233    fn is_empty(&self) -> bool {
234        self.lines.is_empty()
235    }
236}
237
238// Helper function to format SystemTime
239fn format_system_time(time: SystemTime) -> String {
240    let datetime: DateTime<Local> = time.into();
241    datetime.format("%Y-%m-%d %H:%M:%S").to_string()
242}
243
244impl fmt::Display for Aggregator {
245    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
246        // Format the start time
247        let start_time_str = format_system_time(self.start_time);
248
249        // Get and format the current time
250        let current_time = std::time::SystemTime::now();
251        let end_time_str = format_system_time(current_time);
252
253        // Write the header with formatted time window
254        writeln!(
255            f,
256            "\x1b[36m>>> Aggregated Logs ({}) >>>\x1b[0m",
257            start_time_str
258        )?;
259
260        // Write each log line
261        for line in self.lines.iter() {
262            writeln!(f, "{}", line)?;
263        }
264        writeln!(
265            f,
266            "\x1b[36m<<< Aggregated Logs ({}) <<<\x1b[0m",
267            end_time_str
268        )?;
269        Ok(())
270    }
271}
272
273/// Messages that can be sent to the LogClientActor remotely.
274#[derive(
275    Debug,
276    Clone,
277    Serialize,
278    Deserialize,
279    Named,
280    Handler,
281    HandleClient,
282    RefClient
283)]
284pub enum LogMessage {
285    /// Log details
286    Log {
287        /// The hostname of the process that generated the log
288        hostname: String,
289        /// String representation of the ProcAddr that generated the log
290        proc_id: String,
291        /// The target output stream (stdout or stderr)
292        output_target: OutputTarget,
293        /// The log payload as bytes
294        payload: wirevalue::Any,
295    },
296
297    /// Flush the log
298    Flush {
299        /// Indicate if the current flush is synced or non-synced.
300        /// If synced, a version number is available. Otherwise, none.
301        sync_version: Option<u64>,
302    },
303}
304
305/// Messages that can be sent to the LogClient locally.
306#[derive(
307    Debug,
308    Clone,
309    Serialize,
310    Deserialize,
311    Named,
312    Handler,
313    HandleClient,
314    RefClient
315)]
316#[expect(
317    clippy::large_enum_variant,
318    reason = "actor message enum with Handler/HandleClient/RefClient derives; boxing fields ripples into client/handler call sites and may require derive-macro changes — separate diff"
319)]
320pub enum LogClientMessage {
321    SetAggregate {
322        /// The time window in seconds to aggregate logs. If None, aggregation is disabled.
323        aggregate_window_sec: Option<u64>,
324    },
325
326    /// Synchronously flush all the logs from all the procs. This is for client to call.
327    StartSyncFlush {
328        /// Expect these many procs to ack the flush message.
329        expected_procs: usize,
330        /// Return once we have received the acks from all the procs
331        reply: OncePortRef<()>,
332        /// Return to the caller the current flush version
333        version: OncePortRef<u64>,
334    },
335}
336
337/// Trait for sending logs
338#[async_trait]
339pub trait LogSender: Send + Sync {
340    /// Send a log payload in bytes
341    fn send(&mut self, target: OutputTarget, payload: Vec<Vec<u8>>) -> anyhow::Result<()>;
342
343    /// Flush the log channel, ensuring all messages are delivered
344    /// Returns when the flush message has been acknowledged
345    fn flush(&mut self) -> anyhow::Result<()>;
346}
347
348/// Represents the target output stream (stdout or stderr)
349#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
350pub enum OutputTarget {
351    /// Standard output stream
352    Stdout,
353    /// Standard error stream
354    Stderr,
355}
356
357#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
358pub enum Stream {
359    /// Standard output stream
360    ChildStdout,
361    /// Standard error stream
362    ChildStderr,
363}
364
365/// Write the log to a local unix channel so some actors can listen to it and stream the log back.
366pub struct LocalLogSender {
367    hostname: String,
368    proc_id: String,
369    tx: ChannelTx<LogMessage>,
370    status: Receiver<TxStatus>,
371}
372
373impl LocalLogSender {
374    fn new(log_channel: ChannelAddr, proc_id: &ProcAddr) -> Result<Self, anyhow::Error> {
375        let tx = channel::dial::<LogMessage>(log_channel)?;
376        let status = tx.status().clone();
377
378        let hostname = hostname::get()
379            .unwrap_or_else(|_| "unknown_host".into())
380            .into_string()
381            .unwrap_or("unknown_host".to_string());
382        Ok(Self {
383            hostname,
384            proc_id: proc_id.to_string(),
385            tx,
386            status,
387        })
388    }
389}
390
391#[async_trait]
392impl LogSender for LocalLogSender {
393    fn send(&mut self, target: OutputTarget, payload: Vec<Vec<u8>>) -> anyhow::Result<()> {
394        if TxStatus::Active == *self.status.borrow() {
395            self.tx.post(LogMessage::Log {
396                hostname: self.hostname.clone(),
397                proc_id: self.proc_id.clone(),
398                output_target: target,
399                payload: wirevalue::Any::serialize(&payload)?,
400            });
401        }
402
403        Ok(())
404    }
405
406    fn flush(&mut self) -> anyhow::Result<()> {
407        // send will make sure message is delivered
408        if TxStatus::Active == *self.status.borrow() {
409            self.tx.post(LogMessage::Flush { sync_version: None });
410        }
411        Ok(())
412    }
413}
414
415/// Message sent to FileMonitor
416#[derive(Debug, Clone, Serialize, Deserialize, Named)]
417pub struct FileMonitorMessage {
418    lines: Vec<String>,
419}
420wirevalue::register_type!(FileMonitorMessage);
421
422/// File appender, coordinates write access to a file via a channel.
423pub struct FileAppender {
424    stdout_addr: ChannelAddr,
425    stderr_addr: ChannelAddr,
426    #[allow(dead_code)] // Tasks are self terminating
427    stdout_task: JoinHandle<()>,
428    #[allow(dead_code)]
429    stderr_task: JoinHandle<()>,
430    stop: Arc<Notify>,
431}
432
433impl fmt::Debug for FileAppender {
434    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
435        f.debug_struct("FileMonitor")
436            .field("stdout_addr", &self.stdout_addr)
437            .field("stderr_addr", &self.stderr_addr)
438            .finish()
439    }
440}
441
442impl FileAppender {
443    /// Create a new FileAppender with aggregated log files for stdout and stderr
444    /// Returns None if file creation fails
445    pub fn new() -> Option<Self> {
446        let stop = Arc::new(Notify::new());
447        // TODO make it configurable
448        let file_name_tag = hostname::get()
449            .unwrap_or_else(|_| "unknown_host".into())
450            .into_string()
451            .unwrap_or("unknown_host".to_string());
452
453        // Create stdout file and task
454        let (stdout_path, stdout_writer) =
455            match get_unique_local_log_destination(&file_name_tag, OutputTarget::Stdout) {
456                Some(writer) => writer,
457                None => {
458                    tracing::warn!("failed to create stdout file");
459                    return None;
460                }
461            };
462        let (stdout_addr, stdout_rx) = {
463            let _guard = tracing::span!(Level::INFO, "appender", file = "stdout").entered();
464            match channel::serve(ChannelAddr::any(ChannelTransport::Unix)) {
465                Ok((addr, rx)) => (addr, rx),
466                Err(e) => {
467                    tracing::warn!("failed to serve stdout channel: {}", e);
468                    return None;
469                }
470            }
471        };
472        let stdout_stop = stop.clone();
473        let stdout_task = tokio::spawn(file_monitor_task(
474            stdout_rx,
475            stdout_writer,
476            OutputTarget::Stdout,
477            stdout_stop,
478        ));
479
480        // Create stderr file and task
481        let (stderr_path, stderr_writer) =
482            match get_unique_local_log_destination(&file_name_tag, OutputTarget::Stderr) {
483                Some(writer) => writer,
484                None => {
485                    tracing::warn!("failed to create stderr file");
486                    return None;
487                }
488            };
489        let (stderr_addr, stderr_rx) = {
490            let _guard = tracing::span!(Level::INFO, "appender", file = "stderr").entered();
491            match channel::serve(ChannelAddr::any(ChannelTransport::Unix)) {
492                Ok((addr, rx)) => (addr, rx),
493                Err(e) => {
494                    tracing::warn!("failed to serve stderr channel: {}", e);
495                    return None;
496                }
497            }
498        };
499        let stderr_stop = stop.clone();
500        let stderr_task = tokio::spawn(file_monitor_task(
501            stderr_rx,
502            stderr_writer,
503            OutputTarget::Stderr,
504            stderr_stop,
505        ));
506
507        tracing::debug!(
508            "FileAppender: created for stdout {} stderr {} ",
509            stdout_path.display(),
510            stderr_path.display()
511        );
512
513        Some(Self {
514            stdout_addr,
515            stderr_addr,
516            stdout_task,
517            stderr_task,
518            stop,
519        })
520    }
521
522    /// Get a channel address for the specified output target
523    pub fn addr_for(&self, target: OutputTarget) -> ChannelAddr {
524        match target {
525            OutputTarget::Stdout => self.stdout_addr.clone(),
526            OutputTarget::Stderr => self.stderr_addr.clone(),
527        }
528    }
529}
530
531impl Drop for FileAppender {
532    fn drop(&mut self) {
533        // Trigger stop signal to notify tasks to exit
534        self.stop.notify_waiters();
535        tracing::debug!("FileMonitor: dropping, stop signal sent, tasks will flush and exit");
536    }
537}
538
539/// Task that receives lines from StreamFwds and writes them to the aggregated file
540async fn file_monitor_task(
541    mut rx: ChannelRx<FileMonitorMessage>,
542    mut writer: Box<dyn io::AsyncWrite + Send + Unpin + 'static>,
543    target: OutputTarget,
544    stop: Arc<Notify>,
545) {
546    loop {
547        tokio::select! {
548            msg = rx.recv() => {
549                match msg {
550                    Ok(msg) => {
551                        // Write lines to aggregated file
552                        for line in &msg.lines {
553                            if let Err(e) = writer.write_all(line.as_bytes()).await {
554                                tracing::warn!("FileMonitor: failed to write line to file: {}", e);
555                                continue;
556                            }
557                            if let Err(e) = writer.write_all(b"\n").await {
558                                tracing::warn!("FileMonitor: failed to write newline to file: {}", e);
559                            }
560                        }
561                        if let Err(e) = writer.flush().await {
562                            tracing::warn!("FileMonitor: failed to flush file: {}", e);
563                        }
564                    }
565                    Err(e) => {
566                        // Channel error
567                        tracing::debug!("FileMonitor task for {:?}: channel error: {}", target, e);
568                        break;
569                    }
570                }
571            }
572            _ = stop.notified() => {
573                tracing::debug!("FileMonitor task for {:?}: stop signal received", target);
574                break;
575            }
576        }
577    }
578
579    // Graceful shutdown: flush one last time
580    if let Err(e) = writer.flush().await {
581        tracing::warn!("FileMonitor: failed final flush: {}", e);
582    }
583    tracing::debug!("FileMonitor task for {:?} exiting", target);
584}
585
586fn create_unique_file_writer(
587    file_name_tag: &str,
588    output_target: OutputTarget,
589    env: env::Env,
590) -> Result<(PathBuf, Box<dyn io::AsyncWrite + Send + Unpin + 'static>)> {
591    let suffix = match output_target {
592        OutputTarget::Stderr => "stderr",
593        OutputTarget::Stdout => "stdout",
594    };
595    let (path, filename) = log_file_path(env, None)?;
596    let path = Path::new(&path);
597    let mut full_path = PathBuf::from(path);
598
599    let uuid = ShortUuid::generate();
600
601    full_path.push(format!(
602        "{}_{}_{}.{}",
603        filename, file_name_tag, uuid, suffix
604    ));
605    let file = std::fs::OpenOptions::new()
606        .create(true)
607        .append(true)
608        .open(full_path.clone())?;
609    let tokio_file = tokio::fs::File::from_std(file);
610    // TODO: should we buffer this?
611    Ok((full_path, Box::new(tokio_file)))
612}
613
614fn get_unique_local_log_destination(
615    file_name_tag: &str,
616    output_target: OutputTarget,
617) -> Option<(PathBuf, Box<dyn io::AsyncWrite + Send + Unpin + 'static>)> {
618    let env: env::Env = env::Env::current();
619    if env == env::Env::Local && !hyperactor_config::global::get(FORCE_FILE_LOG) {
620        tracing::debug!("not creating log file because of env type");
621        None
622    } else {
623        match create_unique_file_writer(file_name_tag, output_target, env) {
624            Ok((a, b)) => Some((a, b)),
625            Err(e) => {
626                tracing::warn!("failed to create unique file writer: {}", e);
627                None
628            }
629        }
630    }
631}
632
633/// Create a writer for stdout or stderr
634fn std_writer(target: OutputTarget) -> Box<dyn io::AsyncWrite + Send + Unpin> {
635    // Return the appropriate standard output or error writer
636    match target {
637        OutputTarget::Stdout => Box::new(tokio::io::stdout()),
638        OutputTarget::Stderr => Box::new(tokio::io::stderr()),
639    }
640}
641
642/// Copy bytes from `reader` to `writer`, forward to log_sender, and forward to FileMonitor.
643/// The same formatted lines go to both log_sender and file_monitor.
644async fn tee(
645    mut reader: impl AsyncRead + Unpin + Send + 'static,
646    mut std_writer: Box<dyn io::AsyncWrite + Send + Unpin>,
647    log_sender: Option<Box<dyn LogSender + Send>>,
648    file_monitor_addr: Option<ChannelAddr>,
649    target: OutputTarget,
650    prefix: Option<String>,
651    stop: Arc<Notify>,
652    recent_lines_buf: RotatingLineBuffer,
653) -> Result<(), io::Error> {
654    let mut buf = [0u8; 8192];
655    let mut line_buffer = Vec::with_capacity(MAX_LINE_SIZE);
656    let mut log_sender = log_sender;
657
658    // Dial the file monitor channel if provided
659    let mut file_monitor_tx: Option<ChannelTx<FileMonitorMessage>> =
660        file_monitor_addr.and_then(|addr| match channel::dial(addr.clone()) {
661            Ok(tx) => Some(tx),
662            Err(e) => {
663                tracing::warn!("Failed to dial file monitor channel {}: {}", addr, e);
664                None
665            }
666        });
667
668    loop {
669        tokio::select! {
670            read_result = reader.read(&mut buf) => {
671                match read_result {
672                    Ok(n) => {
673                        if n == 0 {
674                            // EOF reached
675                            tracing::debug!("EOF reached in tee");
676                            break;
677                        }
678
679                        // Write to console
680                        if let Err(e) = std_writer.write_all(&buf[..n]).await {
681                            tracing::warn!("error writing to std: {}", e);
682                        }
683
684                        // Process bytes into lines for log_sender and FileMonitor
685                        let mut completed_lines = Vec::new();
686
687                        for &byte in &buf[..n] {
688                            if byte == b'\n' {
689                                // Complete line found
690                                let mut line = String::from_utf8_lossy(&line_buffer).to_string();
691
692                                // Truncate if too long, respecting UTF-8 boundaries
693                                // (multi-byte chars like emojis can be up to 4 bytes)
694                                if line.len() > MAX_LINE_SIZE {
695                                    let mut truncate_at = MAX_LINE_SIZE;
696                                    while truncate_at > 0 && !line.is_char_boundary(truncate_at) {
697                                        truncate_at -= 1;
698                                    }
699                                    line.truncate(truncate_at);
700                                    line.push_str("... [TRUNCATED]");
701                                }
702
703                                // Prepend with prefix if configured
704                                let final_line = if let Some(ref p) = prefix {
705                                    format!("[{}] {}", p, line)
706                                } else {
707                                    line
708                                };
709
710                                completed_lines.push(final_line);
711                                line_buffer.clear();
712                            } else {
713                                line_buffer.push(byte);
714                            }
715                        }
716
717                        // Send completed lines to both log_sender and FileAppender
718                        if !completed_lines.is_empty() {
719                            if let Some(ref mut sender) = log_sender {
720                                let bytes: Vec<Vec<u8>> = completed_lines.iter()
721                                    .map(|s| s.as_bytes().to_vec())
722                                    .collect();
723                                if let Err(e) = sender.send(target, bytes) {
724                                    tracing::warn!("error sending to log_sender: {}", e);
725                                }
726                            }
727
728                            // Send to FileMonitor via hyperactor channel
729                            if let Some(ref mut tx) = file_monitor_tx {
730                                let msg = FileMonitorMessage {
731                                    lines: completed_lines,
732                                };
733                                // Use post() to avoid blocking
734                                tx.post(msg);
735                            }
736                        }
737
738                        recent_lines_buf.try_add_data(&buf, n);
739                    },
740                    Err(e) => {
741                        tracing::debug!("read error in tee: {}", e);
742                        return Err(e);
743                    }
744                }
745            },
746            _ = stop.notified() => {
747                tracing::debug!("stop signal received in tee");
748                break;
749            }
750        }
751    }
752
753    std_writer.flush().await?;
754
755    // Send any remaining partial line
756    if !line_buffer.is_empty() {
757        let mut line = String::from_utf8_lossy(&line_buffer).to_string();
758        // Truncate if too long, respecting UTF-8 boundaries
759        // (multi-byte chars like emojis can be up to 4 bytes)
760        if line.len() > MAX_LINE_SIZE {
761            let mut truncate_at = MAX_LINE_SIZE;
762            while truncate_at > 0 && !line.is_char_boundary(truncate_at) {
763                truncate_at -= 1;
764            }
765            line.truncate(truncate_at);
766            line.push_str("... [TRUNCATED]");
767        }
768        let final_line = if let Some(ref p) = prefix {
769            format!("[{}] {}", p, line)
770        } else {
771            line
772        };
773
774        let final_lines = vec![final_line];
775
776        // Send to log_sender
777        if let Some(ref mut sender) = log_sender {
778            let bytes: Vec<Vec<u8>> = final_lines.iter().map(|s| s.as_bytes().to_vec()).collect();
779            let _ = sender.send(target, bytes);
780        }
781
782        // Send to FileMonitor
783        if let Some(ref mut tx) = file_monitor_tx {
784            let msg = FileMonitorMessage { lines: final_lines };
785            tx.post(msg);
786        }
787    }
788
789    // Flush log_sender
790    if let Some(ref mut sender) = log_sender {
791        let _ = sender.flush();
792    }
793
794    Ok(())
795}
796
797#[derive(Debug, Clone)]
798struct RotatingLineBuffer {
799    recent_lines: Arc<RwLock<VecDeque<String>>>,
800    max_buffer_size: usize,
801}
802
803impl RotatingLineBuffer {
804    fn try_add_data(&self, buf: &[u8], buf_end: usize) {
805        let data_str = String::from_utf8_lossy(&buf[..buf_end]);
806        let lines: Vec<&str> = data_str.lines().collect();
807
808        if let Ok(mut recent_lines_guard) = self.recent_lines.try_write() {
809            for line in lines {
810                if !line.is_empty() {
811                    recent_lines_guard.push_back(line.to_string());
812                    if recent_lines_guard.len() > self.max_buffer_size {
813                        recent_lines_guard.pop_front();
814                    }
815                }
816            }
817        } else {
818            tracing::debug!("Failed to acquire write lock on recent_lines buffer in tee");
819        }
820    }
821
822    async fn peek(&self) -> Vec<String> {
823        let lines = self.recent_lines.read().await;
824        let start_idx = if lines.len() > self.max_buffer_size {
825            lines.len() - self.max_buffer_size
826        } else {
827            0
828        };
829
830        lines.range(start_idx..).cloned().collect()
831    }
832}
833
834/// Given a stream forwards data to the provided channel.
835pub struct StreamFwder {
836    teer: JoinHandle<Result<(), io::Error>>,
837    // Shared buffer for peek functionality
838    recent_lines_buf: RotatingLineBuffer,
839    // Shutdown signal to stop the monitoring loop
840    stop: Arc<Notify>,
841}
842
843impl StreamFwder {
844    /// Create a new StreamFwder instance, and start monitoring the provided path.
845    /// Once started Monitor will
846    /// - forward logs to log_sender
847    /// - forward logs to file_monitor (if available)
848    /// - pipe reader to target
849    /// - And capture last `max_buffer_size` which can be used to inspect file contents via `peek`.
850    pub fn start(
851        reader: impl AsyncRead + Unpin + Send + 'static,
852        file_monitor_addr: Option<ChannelAddr>,
853        target: OutputTarget,
854        max_buffer_size: usize,
855        log_channel: Option<ChannelAddr>,
856        proc_id: &ProcAddr,
857        local_rank: usize,
858    ) -> Self {
859        let prefix = match hyperactor_config::global::get(PREFIX_WITH_RANK) {
860            true => Some(local_rank.to_string()),
861            false => None,
862        };
863        let std_writer = std_writer(target);
864
865        Self::start_with_writer(
866            reader,
867            std_writer,
868            file_monitor_addr,
869            target,
870            max_buffer_size,
871            log_channel,
872            proc_id,
873            prefix,
874        )
875    }
876
877    /// Create a new StreamFwder instance with a custom writer (used in tests).
878    fn start_with_writer(
879        reader: impl AsyncRead + Unpin + Send + 'static,
880        std_writer: Box<dyn io::AsyncWrite + Send + Unpin>,
881        file_monitor_addr: Option<ChannelAddr>,
882        target: OutputTarget,
883        max_buffer_size: usize,
884        log_channel: Option<ChannelAddr>,
885        proc_id: &ProcAddr,
886        prefix: Option<String>,
887    ) -> Self {
888        // Sanity: when there is no file sink, no log forwarding, and
889        // `tail_size == 0`, the child should have **inherited** stdio
890        // and no `StreamFwder` should exist. In that case console
891        // mirroring happens via inheritance, not via `StreamFwder`.
892        // If we hit this, we piped unnecessarily.
893        debug_assert!(
894            file_monitor_addr.is_some() || max_buffer_size > 0 || log_channel.is_some(),
895            "StreamFwder started with no sinks and no tail"
896        );
897
898        let stop = Arc::new(Notify::new());
899        let recent_lines_buf = RotatingLineBuffer {
900            recent_lines: Arc::new(RwLock::new(VecDeque::<String>::with_capacity(
901                max_buffer_size,
902            ))),
903            max_buffer_size,
904        };
905
906        let log_sender: Option<Box<dyn LogSender + Send>> = if let Some(addr) = log_channel {
907            match LocalLogSender::new(addr, proc_id) {
908                Ok(s) => Some(Box::new(s) as Box<dyn LogSender + Send>),
909                Err(e) => {
910                    tracing::error!("failed to create log sender: {}", e);
911                    None
912                }
913            }
914        } else {
915            None
916        };
917
918        let teer_stop = stop.clone();
919        let recent_line_buf_clone = recent_lines_buf.clone();
920        let teer = tokio::spawn(async move {
921            tee(
922                reader,
923                std_writer,
924                log_sender,
925                file_monitor_addr,
926                target,
927                prefix,
928                teer_stop,
929                recent_line_buf_clone,
930            )
931            .await
932        });
933
934        StreamFwder {
935            teer,
936            recent_lines_buf,
937            stop,
938        }
939    }
940
941    pub async fn abort(self) -> (Vec<String>, Result<(), anyhow::Error>) {
942        self.stop.notify_waiters();
943
944        let lines = self.peek().await;
945        let teer_result = self.teer.await;
946
947        let result: Result<(), anyhow::Error> = match teer_result {
948            Ok(inner) => inner.map_err(anyhow::Error::from),
949            Err(e) => Err(e.into()),
950        };
951
952        (lines, result)
953    }
954
955    /// Inspect the latest `max_buffer` lines read from the file being monitored
956    /// Returns lines in chronological order (oldest first)
957    pub async fn peek(&self) -> Vec<String> {
958        self.recent_lines_buf.peek().await
959    }
960}
961
962/// Messages that can be sent to the LogForwarder
963#[derive(
964    Debug,
965    Clone,
966    Serialize,
967    Deserialize,
968    Named,
969    Handler,
970    HandleClient,
971    RefClient
972)]
973pub enum LogForwardMessage {
974    /// Receive the log from the parent process and forward it to the client.
975    Forward {},
976
977    /// If to stream the log back to the client.
978    SetMode { stream_to_client: bool },
979
980    /// Flush the log with a version number.
981    ForceSyncFlush { version: u64 },
982}
983
984/// A log forwarder that receives the log from its parent process and forward it back to the client
985#[hyperactor::export(LogForwardMessage)]
986#[hyperactor::spawnable]
987pub struct LogForwardActor {
988    rx: ChannelRx<LogMessage>,
989    flush_tx: Arc<Mutex<ChannelTx<LogMessage>>>,
990    next_flush_deadline: SystemTime,
991    logging_client_ref: ActorRef<LogClientActor>,
992    stream_to_client: bool,
993}
994
995#[async_trait]
996impl Actor for LogForwardActor {
997    async fn init(&mut self, this: &Instance<Self>) -> Result<(), anyhow::Error> {
998        this.set_system();
999        this.post_after(this, LogForwardMessage::Forward {}, Duration::from_secs(0));
1000
1001        // Make sure we start the flush loop periodically so the log channel will not deadlock.
1002        self.flush_tx
1003            .lock()
1004            .await
1005            .send(LogMessage::Flush { sync_version: None })
1006            .await?;
1007        Ok(())
1008    }
1009}
1010
1011#[async_trait]
1012impl hyperactor::RemoteSpawn for LogForwardActor {
1013    type Params = ActorRef<LogClientActor>;
1014
1015    async fn new(
1016        logging_client_ref: Self::Params,
1017        _environment: &ActorEnvironment,
1018    ) -> Result<Self> {
1019        let log_channel: ChannelAddr = match std::env::var(BOOTSTRAP_LOG_CHANNEL) {
1020            Ok(channel) => channel.parse()?,
1021            Err(err) => {
1022                tracing::debug!(
1023                    "log forwarder actor failed to read env var {}: {}",
1024                    BOOTSTRAP_LOG_CHANNEL,
1025                    err
1026                );
1027                // TODO: an empty channel to serve
1028                ChannelAddr::any(ChannelTransport::Unix)
1029            }
1030        };
1031        tracing::info!(
1032            "log forwarder {} serve at {}",
1033            std::process::id(),
1034            log_channel
1035        );
1036
1037        let rx = match channel::serve(log_channel.clone()) {
1038            Ok((_, rx)) => rx,
1039            Err(err) => {
1040                // This can happen if we are not spanwed on a separate process like local.
1041                // For local mesh, log streaming anyway is not needed.
1042                tracing::error!(
1043                    "log forwarder actor failed to bootstrap on given channel {}: {}",
1044                    log_channel,
1045                    err
1046                );
1047                channel::serve(ChannelAddr::any(ChannelTransport::Unix))?.1
1048            }
1049        };
1050
1051        // Dial the same channel to send flush message to drain the log queue.
1052        let flush_tx = Arc::new(Mutex::new(channel::dial::<LogMessage>(log_channel)?));
1053        let now = std::time::SystemTime::now();
1054
1055        Ok(Self {
1056            rx,
1057            flush_tx,
1058            next_flush_deadline: now,
1059            logging_client_ref,
1060            stream_to_client: true,
1061        })
1062    }
1063}
1064
1065#[async_trait]
1066#[hyperactor::handle(LogForwardMessage)]
1067impl LogForwardMessageHandler for LogForwardActor {
1068    async fn forward(&mut self, ctx: &Context<Self>) -> Result<(), anyhow::Error> {
1069        match self.rx.recv().await {
1070            Ok(LogMessage::Flush { sync_version }) => {
1071                let now = std::time::SystemTime::now();
1072                match sync_version {
1073                    None => {
1074                        // Schedule another flush to keep the log channel from deadlocking.
1075                        let delay = Duration::from_secs(1);
1076                        if now >= self.next_flush_deadline {
1077                            self.next_flush_deadline = now + delay;
1078                            let flush_tx = self.flush_tx.clone();
1079                            tokio::spawn(async move {
1080                                tokio::time::sleep(delay).await;
1081                                if let Err(e) = flush_tx
1082                                    .lock()
1083                                    .await
1084                                    .send(LogMessage::Flush { sync_version: None })
1085                                    .await
1086                                {
1087                                    tracing::error!("failed to send flush message: {}", e);
1088                                }
1089                            });
1090                        }
1091                    }
1092                    version => {
1093                        self.logging_client_ref.flush(ctx, version).await?;
1094                    }
1095                }
1096            }
1097            Ok(LogMessage::Log {
1098                hostname,
1099                proc_id,
1100                output_target,
1101                payload,
1102            }) => {
1103                if self.stream_to_client {
1104                    self.logging_client_ref
1105                        .log(ctx, hostname, proc_id, output_target, payload)
1106                        .await?;
1107                }
1108            }
1109            Err(e) => {
1110                return Err(e.into());
1111            }
1112        }
1113
1114        // This is not ideal as we are using raw tx/rx.
1115        ctx.post_after(ctx, LogForwardMessage::Forward {}, Duration::from_secs(0));
1116
1117        Ok(())
1118    }
1119
1120    async fn set_mode(
1121        &mut self,
1122        _ctx: &Context<Self>,
1123        stream_to_client: bool,
1124    ) -> Result<(), anyhow::Error> {
1125        self.stream_to_client = stream_to_client;
1126        Ok(())
1127    }
1128
1129    async fn force_sync_flush(
1130        &mut self,
1131        _cx: &Context<Self>,
1132        version: u64,
1133    ) -> Result<(), anyhow::Error> {
1134        self.flush_tx
1135            .lock()
1136            .await
1137            .send(LogMessage::Flush {
1138                sync_version: Some(version),
1139            })
1140            .await
1141            .map_err(anyhow::Error::from)
1142    }
1143}
1144
1145/// Deserialize a serialized message and split it into UTF-8 lines
1146fn deserialize_message_lines(serialized_message: &wirevalue::Any) -> Result<Vec<Vec<String>>> {
1147    // Try to deserialize as Vec<Vec<u8>> first (multiple byte arrays)
1148    if let Ok(message_bytes) = serialized_message.deserialized::<Vec<Vec<u8>>>() {
1149        let mut result = Vec::new();
1150        for bytes in message_bytes {
1151            let message_str = String::from_utf8(bytes)?;
1152            let lines: Vec<String> = message_str.lines().map(|s| s.to_string()).collect();
1153            result.push(lines);
1154        }
1155        return Ok(result);
1156    }
1157
1158    // If that fails, try to deserialize as String and wrap in Vec<Vec<String>>
1159    if let Ok(message) = serialized_message.deserialized::<String>() {
1160        let lines: Vec<String> = message.lines().map(|s| s.to_string()).collect();
1161        return Ok(vec![lines]);
1162    }
1163
1164    // If both fail, return an error
1165    anyhow::bail!("failed to deserialize message as either Vec<Vec<u8>> or String")
1166}
1167
1168/// A client to receive logs from remote processes
1169#[derive(Debug)]
1170#[hyperactor::export(LogMessage, LogClientMessage)]
1171#[hyperactor::spawnable]
1172pub struct LogClientActor {
1173    aggregate_window_sec: Option<u64>,
1174    aggregators: HashMap<OutputTarget, Aggregator>,
1175    last_flush_time: SystemTime,
1176    next_flush_deadline: Option<SystemTime>,
1177
1178    // For flush sync barrier
1179    current_flush_version: u64,
1180    current_flush_port: Option<OncePortRef<()>>,
1181    current_unflushed_procs: usize,
1182}
1183
1184impl Default for LogClientActor {
1185    fn default() -> Self {
1186        // Initialize aggregators
1187        let mut aggregators = HashMap::new();
1188        aggregators.insert(OutputTarget::Stderr, Aggregator::new());
1189        aggregators.insert(OutputTarget::Stdout, Aggregator::new());
1190
1191        Self {
1192            aggregate_window_sec: Some(DEFAULT_AGGREGATE_WINDOW_SEC),
1193            aggregators,
1194            last_flush_time: std::time::SystemTime::now(),
1195            next_flush_deadline: None,
1196            current_flush_version: 0,
1197            current_flush_port: None,
1198            current_unflushed_procs: 0,
1199        }
1200    }
1201}
1202
1203#[async_trait]
1204impl Actor for LogClientActor {
1205    async fn init(&mut self, this: &Instance<Self>) -> Result<(), anyhow::Error> {
1206        this.set_system();
1207        Ok(())
1208    }
1209}
1210
1211impl LogClientActor {
1212    fn print_aggregators(&mut self) {
1213        for (output_target, aggregator) in self.aggregators.iter_mut() {
1214            if aggregator.is_empty() {
1215                continue;
1216            }
1217            match output_target {
1218                OutputTarget::Stdout => {
1219                    println!("{}", aggregator);
1220                }
1221                OutputTarget::Stderr => {
1222                    eprintln!("{}", aggregator);
1223                }
1224            }
1225
1226            // Reset the aggregator
1227            aggregator.reset();
1228        }
1229    }
1230
1231    fn print_log_line(hostname: &str, proc_id: &str, output_target: OutputTarget, line: String) {
1232        let message = format!("[{} {}] {}", hostname, proc_id, line);
1233
1234        #[cfg(test)]
1235        crate::logging::test_tap::push(&message);
1236
1237        match output_target {
1238            OutputTarget::Stdout => println!("{}", message),
1239            OutputTarget::Stderr => eprintln!("{}", message),
1240        }
1241    }
1242
1243    fn flush_internal(&mut self) {
1244        self.print_aggregators();
1245        self.last_flush_time = std::time::SystemTime::now();
1246        self.next_flush_deadline = None;
1247    }
1248}
1249
1250impl Drop for LogClientActor {
1251    fn drop(&mut self) {
1252        // Flush the remaining logs before shutting down
1253        self.print_aggregators();
1254    }
1255}
1256
1257#[async_trait]
1258#[hyperactor::handle(LogMessage)]
1259impl LogMessageHandler for LogClientActor {
1260    async fn log(
1261        &mut self,
1262        cx: &Context<Self>,
1263        hostname: String,
1264        proc_id: String,
1265        output_target: OutputTarget,
1266        payload: wirevalue::Any,
1267    ) -> Result<(), anyhow::Error> {
1268        // Deserialize the message and process line by line with UTF-8
1269        let message_line_groups = deserialize_message_lines(&payload)?;
1270        let hostname = hostname.as_str();
1271
1272        let message_lines: Vec<String> = message_line_groups.into_iter().flatten().collect();
1273        match self.aggregate_window_sec {
1274            None => {
1275                for line in message_lines {
1276                    Self::print_log_line(hostname, &proc_id, output_target, line);
1277                }
1278                self.last_flush_time = std::time::SystemTime::now();
1279            }
1280            Some(window) => {
1281                for line in message_lines {
1282                    if let Some(aggregator) = self.aggregators.get_mut(&output_target) {
1283                        if let Err(e) = aggregator.add_line(&line) {
1284                            tracing::error!("error adding log line: {}", e);
1285                            // For the sake of completeness, flush the log lines.
1286                            Self::print_log_line(hostname, &proc_id, output_target, line);
1287                        }
1288                    } else {
1289                        tracing::error!("unknown output target: {:?}", output_target);
1290                        // For the sake of completeness, flush the log lines.
1291                        Self::print_log_line(hostname, &proc_id, output_target, line);
1292                    }
1293                }
1294
1295                let new_deadline = self.last_flush_time + Duration::from_secs(window);
1296                let now = std::time::SystemTime::now();
1297                if new_deadline <= now {
1298                    self.flush_internal();
1299                } else {
1300                    let delay = new_deadline.duration_since(now)?;
1301                    match self.next_flush_deadline {
1302                        None => {
1303                            self.next_flush_deadline = Some(new_deadline);
1304                            cx.post_after(cx, LogMessage::Flush { sync_version: None }, delay);
1305                        }
1306                        Some(deadline) => {
1307                            // Some early log lines have alrady triggered the flush.
1308                            if new_deadline < deadline {
1309                                // This can happen if the user has adjusted the aggregation window.
1310                                self.next_flush_deadline = Some(new_deadline);
1311                                cx.post_after(cx, LogMessage::Flush { sync_version: None }, delay);
1312                            }
1313                        }
1314                    }
1315                }
1316            }
1317        }
1318
1319        Ok(())
1320    }
1321
1322    async fn flush(
1323        &mut self,
1324        cx: &Context<Self>,
1325        sync_version: Option<u64>,
1326    ) -> Result<(), anyhow::Error> {
1327        match sync_version {
1328            None => {
1329                self.flush_internal();
1330            }
1331            Some(version) => {
1332                if version != self.current_flush_version {
1333                    tracing::error!(
1334                        "found mismatched flush versions: got {}, expect {}; this can happen if some previous flush didn't finish fully",
1335                        version,
1336                        self.current_flush_version
1337                    );
1338                    return Ok(());
1339                }
1340
1341                if self.current_unflushed_procs == 0 || self.current_flush_port.is_none() {
1342                    // This is a serious issue; it's better to error out.
1343                    anyhow::bail!("found no ongoing flush request");
1344                }
1345                self.current_unflushed_procs -= 1;
1346
1347                tracing::debug!(
1348                    "ack sync flush: version {}; remaining procs: {}",
1349                    self.current_flush_version,
1350                    self.current_unflushed_procs
1351                );
1352
1353                if self.current_unflushed_procs == 0 {
1354                    self.flush_internal();
1355                    let reply = self.current_flush_port.take().unwrap();
1356                    self.current_flush_port = None;
1357                    reply.post(cx, ());
1358                }
1359            }
1360        }
1361
1362        Ok(())
1363    }
1364}
1365
1366#[async_trait]
1367#[hyperactor::handle(LogClientMessage)]
1368impl LogClientMessageHandler for LogClientActor {
1369    async fn set_aggregate(
1370        &mut self,
1371        _cx: &Context<Self>,
1372        aggregate_window_sec: Option<u64>,
1373    ) -> Result<(), anyhow::Error> {
1374        if self.aggregate_window_sec.is_some() && aggregate_window_sec.is_none() {
1375            // Make sure we flush whatever in the aggregators before disabling aggregation.
1376            self.print_aggregators();
1377        }
1378        self.aggregate_window_sec = aggregate_window_sec;
1379        Ok(())
1380    }
1381
1382    async fn start_sync_flush(
1383        &mut self,
1384        cx: &Context<Self>,
1385        expected_procs_flushed: usize,
1386        reply: OncePortRef<()>,
1387        version: OncePortRef<u64>,
1388    ) -> Result<(), anyhow::Error> {
1389        if self.current_unflushed_procs > 0 || self.current_flush_port.is_some() {
1390            tracing::warn!(
1391                "found unfinished ongoing flush: version {}; {} unflushed procs",
1392                self.current_flush_version,
1393                self.current_unflushed_procs,
1394            );
1395        }
1396
1397        self.current_flush_version += 1;
1398        tracing::debug!(
1399            "start sync flush with version {}",
1400            self.current_flush_version
1401        );
1402        self.current_flush_port = Some(reply.clone());
1403        self.current_unflushed_procs = expected_procs_flushed;
1404        version.post(cx, self.current_flush_version);
1405        Ok(())
1406    }
1407}
1408
1409#[cfg(test)]
1410pub mod test_tap {
1411    use std::sync::Mutex;
1412    use std::sync::OnceLock;
1413
1414    use tokio::sync::mpsc::UnboundedReceiver;
1415    use tokio::sync::mpsc::UnboundedSender;
1416
1417    static TAP: OnceLock<UnboundedSender<String>> = OnceLock::new();
1418    static RX: OnceLock<Mutex<UnboundedReceiver<String>>> = OnceLock::new();
1419
1420    // Called by tests to install the sender.
1421    pub fn install(tx: UnboundedSender<String>) {
1422        let _ = TAP.set(tx);
1423    }
1424
1425    // Called by tests to register the receiver so we can drain later.
1426    pub fn set_receiver(rx: UnboundedReceiver<String>) {
1427        let _ = RX.set(Mutex::new(rx));
1428    }
1429
1430    // Used by LogClientActor (under #[cfg(test)]) to push a line.
1431    pub fn push(s: &str) {
1432        if let Some(tx) = TAP.get() {
1433            let _ = tx.send(s.to_string());
1434        }
1435    }
1436
1437    // Tests call this to collect everything observed so far.
1438    pub fn drain() -> Vec<String> {
1439        let mut out = Vec::new();
1440        if let Some(rx) = RX.get() {
1441            let mut rx = rx.lock().unwrap();
1442            while let Ok(line) = rx.try_recv() {
1443                out.push(line);
1444            }
1445        }
1446        out
1447    }
1448}
1449
1450#[cfg(test)]
1451mod tests {
1452
1453    use std::sync::Arc;
1454    use std::sync::Mutex;
1455
1456    use hyperactor::ProcAddr;
1457    use hyperactor::RemoteSpawn;
1458    use hyperactor::channel;
1459    use hyperactor::channel::ChannelAddr;
1460    use hyperactor::channel::ChannelTx;
1461    use hyperactor::channel::Tx;
1462    use hyperactor::mailbox::BoxedMailboxSender;
1463    use hyperactor::mailbox::DialMailboxRouter;
1464    use hyperactor::mailbox::MailboxServer;
1465    use hyperactor::proc::Proc;
1466    use hyperactor::testing::ids::test_proc_id;
1467    use tokio::io::AsyncSeek;
1468    use tokio::io::AsyncSeekExt;
1469    use tokio::io::AsyncWriteExt;
1470    use tokio::io::SeekFrom;
1471    use tokio::sync::mpsc;
1472
1473    use super::*;
1474
1475    /// Result of processing file content
1476    #[derive(Debug)]
1477    struct FileProcessingResult {
1478        /// Complete lines found during processing
1479        lines: Vec<Vec<u8>>,
1480        /// Updated position in the file after processing
1481        new_position: u64,
1482        /// Any remaining incomplete line data, buffered for subsequent reads
1483        incomplete_line_buffer: Vec<u8>,
1484    }
1485
1486    /// Process new file content from a given position, extracting complete lines
1487    /// This function is extracted to enable easier unit testing without file system dependencies
1488    async fn process_file_content<R: AsyncRead + AsyncSeek + Unpin>(
1489        reader: &mut R,
1490        current_position: u64,
1491        file_size: u64,
1492        existing_line_buffer: Vec<u8>,
1493        max_buffer_size: usize,
1494    ) -> Result<FileProcessingResult> {
1495        // If position equals file size, we're at the end
1496        if current_position == file_size {
1497            return Ok(FileProcessingResult {
1498                lines: Vec::new(),
1499                new_position: current_position,
1500                incomplete_line_buffer: existing_line_buffer,
1501            });
1502        }
1503
1504        // Handle potential file truncation/rotation
1505        let actual_position = if current_position > file_size {
1506            tracing::warn!(
1507                "File appears to have been truncated (position {} > file size {}), resetting to start",
1508                current_position,
1509                file_size
1510            );
1511            reader.seek(SeekFrom::Start(0)).await?;
1512            0
1513        } else {
1514            // current_position < file_size
1515            reader.seek(SeekFrom::Start(current_position)).await?;
1516            current_position
1517        };
1518
1519        let mut buf = vec![0u8; 128 * 1024];
1520        let mut line_buffer = existing_line_buffer;
1521        let mut lines = Vec::with_capacity(max_buffer_size);
1522        let mut processed_bytes = 0u64;
1523
1524        loop {
1525            let bytes_read = reader.read(&mut buf).await?;
1526            if bytes_read == 0 {
1527                break;
1528            }
1529
1530            let chunk = &buf[..bytes_read];
1531
1532            let mut start = 0;
1533            while let Some(newline_pos) = chunk[start..].iter().position(|&b| b == b'\n') {
1534                let absolute_pos = start + newline_pos;
1535
1536                line_buffer.extend_from_slice(&chunk[start..absolute_pos]);
1537
1538                if !line_buffer.is_empty() {
1539                    if line_buffer.len() > MAX_LINE_SIZE {
1540                        line_buffer.truncate(MAX_LINE_SIZE);
1541                        line_buffer.extend_from_slice(b"... [TRUNCATED]");
1542                    }
1543
1544                    let line_data = std::mem::replace(&mut line_buffer, Vec::with_capacity(2048));
1545                    lines.push(line_data);
1546                }
1547
1548                start = absolute_pos + 1;
1549
1550                // Check if we've reached the max buffer size after adding each line
1551                if lines.len() >= max_buffer_size {
1552                    // We've processed up to and including the current newline
1553                    // The new position is where we should start reading next time
1554                    let new_position = actual_position + processed_bytes + start as u64;
1555
1556                    // Don't save remaining data - we'll re-read it from the new position
1557                    return Ok(FileProcessingResult {
1558                        lines,
1559                        new_position,
1560                        incomplete_line_buffer: Vec::new(),
1561                    });
1562                }
1563            }
1564
1565            // Only add bytes to processed_bytes if we've fully processed this chunk
1566            processed_bytes += bytes_read as u64;
1567
1568            if start < chunk.len() {
1569                line_buffer.extend_from_slice(&chunk[start..]);
1570            }
1571        }
1572
1573        let new_position = actual_position + processed_bytes;
1574
1575        Ok(FileProcessingResult {
1576            lines,
1577            new_position,
1578            incomplete_line_buffer: line_buffer,
1579        })
1580    }
1581
1582    #[tokio::test]
1583    async fn test_forwarding_log_to_client() {
1584        // Setup the basics
1585        let router = DialMailboxRouter::new();
1586        let (proc_addr, client_rx) =
1587            channel::serve(ChannelAddr::any(ChannelTransport::Unix)).unwrap();
1588        let proc = Proc::configured(
1589            test_proc_id("client_0"),
1590            BoxedMailboxSender::new(router.clone()),
1591        );
1592        proc.clone().serve(client_rx);
1593        let proc_ref: ProcAddr = test_proc_id("client_0");
1594        router.bind(proc_ref, proc_addr.clone());
1595        let client = proc.client("client");
1596
1597        // Spin up both the forwarder and the client
1598        let log_channel = ChannelAddr::any(ChannelTransport::Unix);
1599        // SAFETY: Unit test
1600        unsafe {
1601            std::env::set_var(BOOTSTRAP_LOG_CHANNEL, log_channel.to_string());
1602        }
1603        let log_client_actor = LogClientActor::new((), &ActorEnvironment::default())
1604            .await
1605            .unwrap();
1606        let log_client: ActorRef<LogClientActor> = proc.spawn(log_client_actor).bind();
1607        let log_forwarder_actor =
1608            LogForwardActor::new(log_client.clone(), &ActorEnvironment::default())
1609                .await
1610                .unwrap();
1611        let log_forwarder: ActorRef<LogForwardActor> = proc.spawn(log_forwarder_actor).bind();
1612
1613        // Write some logs that will not be streamed
1614        let tx: ChannelTx<LogMessage> = channel::dial(log_channel).unwrap();
1615        tx.post(LogMessage::Log {
1616            hostname: "my_host".into(),
1617            proc_id: "test_proc".into(),
1618            output_target: OutputTarget::Stderr,
1619            payload: wirevalue::Any::serialize(&"will not stream".to_string()).unwrap(),
1620        });
1621
1622        // Turn on streaming
1623        log_forwarder.set_mode(&client, true).await.unwrap();
1624        tx.post(LogMessage::Log {
1625            hostname: "my_host".into(),
1626            proc_id: "test_proc".into(),
1627            output_target: OutputTarget::Stderr,
1628            payload: wirevalue::Any::serialize(&"will stream".to_string()).unwrap(),
1629        });
1630
1631        // TODO: it is hard to test out anything meaningful here as the client flushes to stdout.
1632    }
1633
1634    #[test]
1635    fn test_deserialize_message_lines_string() {
1636        // Test deserializing a String message with multiple lines
1637        let message = "Line 1\nLine 2\nLine 3".to_string();
1638        let serialized = wirevalue::Any::serialize(&message).unwrap();
1639
1640        let result = deserialize_message_lines(&serialized).unwrap();
1641        assert_eq!(result, vec![vec!["Line 1", "Line 2", "Line 3"]]);
1642
1643        // Test deserializing a Vec<Vec<u8>> message with UTF-8 content
1644        let message_bytes = vec![
1645            "Hello\nWorld".as_bytes().to_vec(),
1646            "UTF-8 \u{1F980}\nTest".as_bytes().to_vec(),
1647        ];
1648        let serialized = wirevalue::Any::serialize(&message_bytes).unwrap();
1649
1650        let result = deserialize_message_lines(&serialized).unwrap();
1651        assert_eq!(
1652            result,
1653            vec![vec!["Hello", "World"], vec!["UTF-8 \u{1F980}", "Test"]]
1654        );
1655
1656        // Test deserializing a single line message
1657        let message = "Single line message".to_string();
1658        let serialized = wirevalue::Any::serialize(&message).unwrap();
1659
1660        let result = deserialize_message_lines(&serialized).unwrap();
1661
1662        assert_eq!(result, vec![vec!["Single line message"]]);
1663
1664        // Test deserializing an empty lines
1665        let message = "\n\n".to_string();
1666        let serialized = wirevalue::Any::serialize(&message).unwrap();
1667
1668        let result = deserialize_message_lines(&serialized).unwrap();
1669
1670        assert_eq!(result, vec![vec!["", ""]]);
1671
1672        // Test error handling for invalid UTF-8 bytes
1673        let invalid_utf8_bytes = vec![vec![0xFF, 0xFE, 0xFD]]; // Invalid UTF-8 sequence in Vec<Vec<u8>>
1674        let serialized = wirevalue::Any::serialize(&invalid_utf8_bytes).unwrap();
1675
1676        let result = deserialize_message_lines(&serialized);
1677
1678        // The function should fail when trying to convert invalid UTF-8 bytes to String
1679        assert!(
1680            result.is_err(),
1681            "Expected deserialization to fail with invalid UTF-8 bytes"
1682        );
1683    }
1684    #[allow(dead_code)]
1685    struct MockLogSender {
1686        log_sender: mpsc::UnboundedSender<(OutputTarget, String)>, // (output_target, content)
1687        flush_called: Arc<Mutex<bool>>,                            // Track if flush was called
1688    }
1689
1690    impl MockLogSender {
1691        #[allow(dead_code)]
1692        fn new(log_sender: mpsc::UnboundedSender<(OutputTarget, String)>) -> Self {
1693            Self {
1694                log_sender,
1695                flush_called: Arc::new(Mutex::new(false)),
1696            }
1697        }
1698    }
1699
1700    #[async_trait]
1701    impl LogSender for MockLogSender {
1702        fn send(
1703            &mut self,
1704            output_target: OutputTarget,
1705            payload: Vec<Vec<u8>>,
1706        ) -> anyhow::Result<()> {
1707            // For testing purposes, convert to string if it's valid UTF-8
1708            let lines: Vec<String> = payload
1709                .iter()
1710                .map(|b| String::from_utf8_lossy(b).trim_end_matches('\n').to_owned())
1711                .collect();
1712
1713            for line in lines {
1714                self.log_sender
1715                    .send((output_target, line))
1716                    .map_err(|e| anyhow::anyhow!("Failed to send log in test: {}", e))?;
1717            }
1718            Ok(())
1719        }
1720
1721        fn flush(&mut self) -> anyhow::Result<()> {
1722            // Mark that flush was called
1723            let mut flush_called = self.flush_called.lock().unwrap();
1724            *flush_called = true;
1725
1726            // For testing purposes, just return Ok
1727            // In a real implementation, this would wait for all messages to be delivered
1728            Ok(())
1729        }
1730    }
1731
1732    #[test]
1733    fn test_string_similarity() {
1734        // Test exact match
1735        assert_eq!(normalized_edit_distance("hello", "hello"), 0.0);
1736
1737        // Test completely different strings
1738        assert_eq!(normalized_edit_distance("hello", "i'mdiff"), 1.0);
1739
1740        // Test similar strings
1741        assert!(normalized_edit_distance("hello", "helo") < 0.5);
1742        assert!(normalized_edit_distance("hello", "hello!") < 0.5);
1743
1744        // Test empty strings
1745        assert_eq!(normalized_edit_distance("", ""), 0.0);
1746        assert_eq!(normalized_edit_distance("hello", ""), 1.0);
1747    }
1748
1749    #[test]
1750    fn test_add_line_to_empty_aggregator() {
1751        let mut aggregator = Aggregator::new();
1752        let result = aggregator.add_line("ERROR 404 not found");
1753
1754        assert!(result.is_ok());
1755        assert_eq!(aggregator.lines.len(), 1);
1756        assert_eq!(aggregator.lines[0].content, "ERROR 404 not found");
1757        assert_eq!(aggregator.lines[0].count, 1);
1758    }
1759
1760    #[test]
1761    fn test_add_line_merges_with_similar_line() {
1762        let mut aggregator = Aggregator::new_with_threshold(0.2);
1763
1764        // Add first line
1765        aggregator.add_line("ERROR 404 timeout").unwrap();
1766        assert_eq!(aggregator.lines.len(), 1);
1767
1768        // Add second line that should merge (similar enough)
1769        aggregator.add_line("ERROR 500 timeout").unwrap();
1770        assert_eq!(aggregator.lines.len(), 1); // Should still be 1 line after merge
1771        assert_eq!(aggregator.lines[0].count, 2);
1772
1773        // Add third line that's too different
1774        aggregator
1775            .add_line("WARNING database connection failed")
1776            .unwrap();
1777        assert_eq!(aggregator.lines.len(), 2); // Should be 2 lines now
1778
1779        // Add fourth line similar to third
1780        aggregator
1781            .add_line("WARNING database connection timed out")
1782            .unwrap();
1783        assert_eq!(aggregator.lines.len(), 2); // Should still be 2 lines
1784        assert_eq!(aggregator.lines[1].count, 2); // Second group has 2 lines
1785    }
1786
1787    #[test]
1788    fn test_aggregation_of_similar_log_lines() {
1789        let mut aggregator = Aggregator::new_with_threshold(0.2);
1790
1791        // Add the provided log lines with small differences
1792        aggregator.add_line("[1 similar log lines] WARNING <<2025, 2025>> -07-30 <<0, 0>> :41:45,366 conda-unpack-fb:292] Found invalid offsets for share/terminfo/i/ims-ansi, falling back to search/replace to update prefixes for this file.").unwrap();
1793        aggregator.add_line("[1 similar log lines] WARNING <<2025, 2025>> -07-30 <<0, 0>> :41:45,351 conda-unpack-fb:292] Found invalid offsets for lib/pkgconfig/ncursesw.pc, falling back to search/replace to update prefixes for this file.").unwrap();
1794        aggregator.add_line("[1 similar log lines] WARNING <<2025, 2025>> -07-30 <<0, 0>> :41:45,366 conda-unpack-fb:292] Found invalid offsets for share/terminfo/k/kt7, falling back to search/replace to update prefixes for this file.").unwrap();
1795
1796        // Check that we have only one aggregated line due to similarity
1797        assert_eq!(aggregator.lines.len(), 1);
1798
1799        // Check that the count is 3
1800        assert_eq!(aggregator.lines[0].count, 3);
1801    }
1802
1803    #[tokio::test]
1804    async fn test_stream_fwd_creation() {
1805        hyperactor_telemetry::initialize_logging_for_test();
1806
1807        let (mut writer, reader) = tokio::io::duplex(1024);
1808        let (log_channel, mut rx) =
1809            channel::serve::<LogMessage>(ChannelAddr::any(ChannelTransport::Unix)).unwrap();
1810
1811        // Create a temporary file for testing the writer
1812        let temp_file = tempfile::NamedTempFile::new().unwrap();
1813        let temp_path = temp_file.path().to_path_buf();
1814
1815        // Create file writer that writes to the temp file (using tokio for async compatibility)
1816        let file_writer = tokio::fs::OpenOptions::new()
1817            .create(true)
1818            .write(true)
1819            .append(true)
1820            .open(&temp_path)
1821            .await
1822            .unwrap();
1823
1824        // Create FileMonitor and get address for stdout
1825        let file_monitor = FileAppender::new();
1826        let file_monitor_addr = file_monitor
1827            .as_ref()
1828            .map(|fm| fm.addr_for(OutputTarget::Stdout));
1829
1830        let the_test_proc_id = test_proc_id("testproc_0");
1831        let monitor = StreamFwder::start_with_writer(
1832            reader,
1833            Box::new(file_writer),
1834            file_monitor_addr,
1835            OutputTarget::Stdout,
1836            3, // max_buffer_size
1837            Some(log_channel),
1838            &the_test_proc_id,
1839            None, // no prefix
1840        );
1841
1842        // Wait a bit for set up to be done
1843        tokio::time::sleep(Duration::from_millis(500)).await;
1844
1845        // Write initial content through the input writer
1846        writer.write_all(b"Initial log line\n").await.unwrap();
1847        writer.flush().await.unwrap();
1848
1849        // Write more content
1850        for i in 1..=5 {
1851            writer
1852                .write_all(format!("New log line {}\n", i).as_bytes())
1853                .await
1854                .unwrap();
1855        }
1856        writer.flush().await.unwrap();
1857
1858        // Wait a bit for the file to be written and the watcher to detect changes
1859        tokio::time::sleep(Duration::from_millis(500)).await;
1860
1861        // Wait until log sender gets message
1862        let timeout = Duration::from_secs(1);
1863        let _ = tokio::time::timeout(timeout, rx.recv())
1864            .await
1865            .unwrap_or_else(|_| panic!("Did not get log message within {:?}", timeout));
1866
1867        // Wait a bit more for all lines to be processed
1868        tokio::time::sleep(Duration::from_millis(200)).await;
1869
1870        let (recent_lines, _result) = monitor.abort().await;
1871
1872        assert!(
1873            recent_lines.len() >= 3,
1874            "Expected buffer with at least 3 lines, got {} lines: {:?}",
1875            recent_lines.len(),
1876            recent_lines
1877        );
1878
1879        let file_contents = std::fs::read_to_string(&temp_path).unwrap();
1880        assert!(
1881            file_contents.contains("Initial log line"),
1882            "Expected temp file to contain 'Initial log line', got: {:?}",
1883            file_contents
1884        );
1885        assert!(
1886            file_contents.contains("New log line 1"),
1887            "Expected temp file to contain 'New log line 1', got: {:?}",
1888            file_contents
1889        );
1890        assert!(
1891            file_contents.contains("New log line 5"),
1892            "Expected temp file to contain 'New log line 5', got: {:?}",
1893            file_contents
1894        );
1895    }
1896
1897    #[test]
1898    fn test_aggregator_custom_threshold() {
1899        // Test with very strict threshold (0.05)
1900        let mut strict_aggregator = Aggregator::new_with_threshold(0.05);
1901        strict_aggregator.add_line("ERROR 404").unwrap();
1902        strict_aggregator.add_line("ERROR 500").unwrap(); // Should not merge due to strict threshold
1903        assert_eq!(strict_aggregator.lines.len(), 2);
1904
1905        // Test with very lenient threshold (0.8)
1906        let mut lenient_aggregator = Aggregator::new_with_threshold(0.8);
1907        lenient_aggregator.add_line("ERROR 404").unwrap();
1908        lenient_aggregator.add_line("WARNING 200").unwrap(); // Should merge due to lenient threshold
1909        assert_eq!(lenient_aggregator.lines.len(), 1);
1910        assert_eq!(lenient_aggregator.lines[0].count, 2);
1911    }
1912
1913    #[test]
1914    fn test_format_system_time() {
1915        let test_time = SystemTime::UNIX_EPOCH + Duration::from_secs(1609459200); // 2021-01-01 00:00:00 UTC
1916        let formatted = format_system_time(test_time);
1917
1918        // Just verify it's a reasonable format (contains date and time components)
1919        assert!(formatted.contains("-"));
1920        assert!(formatted.contains(":"));
1921        assert!(formatted.len() > 10); // Should be reasonable length
1922    }
1923
1924    #[test]
1925    fn test_aggregator_display_formatting() {
1926        let mut aggregator = Aggregator::new();
1927        aggregator.add_line("Test error message").unwrap();
1928        aggregator.add_line("Test error message").unwrap(); // Should merge
1929
1930        let display_string = format!("{}", aggregator);
1931
1932        // Verify the output contains expected elements
1933        assert!(display_string.contains("Aggregated Logs"));
1934        assert!(display_string.contains("[2 similar log lines]"));
1935        assert!(display_string.contains("Test error message"));
1936        assert!(display_string.contains(">>>") && display_string.contains("<<<"));
1937    }
1938
1939    #[tokio::test]
1940    async fn test_local_log_sender_inactive_status() {
1941        let (log_channel, _) =
1942            channel::serve::<LogMessage>(ChannelAddr::any(ChannelTransport::Unix)).unwrap();
1943        let test_proc_id = test_proc_id("testproc_0");
1944        let mut sender = LocalLogSender::new(log_channel, &test_proc_id).unwrap();
1945
1946        // This test verifies that the sender handles inactive status gracefully
1947        // In a real scenario, the channel would be closed, but for testing we just
1948        // verify the send/flush methods don't panic
1949        let result = sender.send(OutputTarget::Stdout, vec![b"test".to_vec()]);
1950        assert!(result.is_ok());
1951
1952        let result = sender.flush();
1953        assert!(result.is_ok());
1954    }
1955
1956    #[test]
1957    fn test_levenshtein_distance_edge_cases() {
1958        // Test with empty strings
1959        assert_eq!(levenshtein_distance("", ""), 0);
1960        assert_eq!(levenshtein_distance("", "hello"), 5);
1961        assert_eq!(levenshtein_distance("hello", ""), 5);
1962
1963        // Test with identical strings
1964        assert_eq!(levenshtein_distance("hello", "hello"), 0);
1965
1966        // Test with single character differences
1967        assert_eq!(levenshtein_distance("hello", "helo"), 1); // deletion
1968        assert_eq!(levenshtein_distance("helo", "hello"), 1); // insertion
1969        assert_eq!(levenshtein_distance("hello", "hallo"), 1); // substitution
1970
1971        // Test with unicode characters
1972        assert_eq!(levenshtein_distance("café", "cafe"), 1);
1973    }
1974
1975    #[test]
1976    fn test_normalized_edit_distance_edge_cases() {
1977        // Test with empty strings
1978        assert_eq!(normalized_edit_distance("", ""), 0.0);
1979
1980        // Test normalization
1981        assert_eq!(normalized_edit_distance("hello", ""), 1.0);
1982        assert_eq!(normalized_edit_distance("", "hello"), 1.0);
1983
1984        // Test that result is always between 0.0 and 1.0
1985        let distance = normalized_edit_distance("completely", "different");
1986        assert!((0.0..=1.0).contains(&distance));
1987    }
1988
1989    #[tokio::test]
1990    async fn test_deserialize_message_lines_edge_cases() {
1991        // Test with empty string
1992        let empty_message = "".to_string();
1993        let serialized = wirevalue::Any::serialize(&empty_message).unwrap();
1994        let result = deserialize_message_lines(&serialized).unwrap();
1995        assert_eq!(result, vec![vec![] as Vec<String>]);
1996
1997        // Test with trailing newline
1998        let trailing_newline = "line1\nline2\n".to_string();
1999        let serialized = wirevalue::Any::serialize(&trailing_newline).unwrap();
2000        let result = deserialize_message_lines(&serialized).unwrap();
2001        assert_eq!(result, vec![vec!["line1", "line2"]]);
2002    }
2003
2004    #[test]
2005    fn test_output_target_serialization() {
2006        // Test that OutputTarget can be serialized and deserialized
2007        let stdout_serialized = serde_json::to_string(&OutputTarget::Stdout).unwrap();
2008        let stderr_serialized = serde_json::to_string(&OutputTarget::Stderr).unwrap();
2009
2010        let stdout_deserialized: OutputTarget = serde_json::from_str(&stdout_serialized).unwrap();
2011        let stderr_deserialized: OutputTarget = serde_json::from_str(&stderr_serialized).unwrap();
2012
2013        assert_eq!(stdout_deserialized, OutputTarget::Stdout);
2014        assert_eq!(stderr_deserialized, OutputTarget::Stderr);
2015    }
2016
2017    #[test]
2018    fn test_log_line_display_formatting() {
2019        let log_line = LogLine::new("Test message".to_string());
2020        let display_string = format!("{}", log_line);
2021
2022        assert!(display_string.contains("[1 similar log lines]"));
2023        assert!(display_string.contains("Test message"));
2024
2025        // Test with higher count
2026        let mut log_line_multi = LogLine::new("Test message".to_string());
2027        log_line_multi.count = 5;
2028        let display_string_multi = format!("{}", log_line_multi);
2029
2030        assert!(display_string_multi.contains("[5 similar log lines]"));
2031        assert!(display_string_multi.contains("Test message"));
2032    }
2033
2034    // Mock reader for testing process_file_content using std::io::Cursor
2035    fn create_mock_reader(data: Vec<u8>) -> std::io::Cursor<Vec<u8>> {
2036        std::io::Cursor::new(data)
2037    }
2038
2039    #[tokio::test]
2040    async fn test_process_file_content_basic() {
2041        let data = b"line1\nline2\nline3\n".to_vec();
2042        let mut reader = create_mock_reader(data.clone());
2043        let max_buf_size = 10;
2044
2045        let result =
2046            process_file_content(&mut reader, 0, data.len() as u64, Vec::new(), max_buf_size)
2047                .await
2048                .unwrap();
2049
2050        assert_eq!(result.lines.len(), 3);
2051        assert_eq!(result.lines[0], b"line1");
2052        assert_eq!(result.lines[1], b"line2");
2053        assert_eq!(result.lines[2], b"line3");
2054        assert_eq!(result.new_position, data.len() as u64);
2055        assert!(result.incomplete_line_buffer.is_empty());
2056    }
2057
2058    #[tokio::test]
2059    async fn test_process_file_content_incomplete_line() {
2060        let data = b"line1\nline2\npartial".to_vec();
2061        let mut reader = create_mock_reader(data.clone());
2062        let max_buf_size = 10;
2063
2064        let result =
2065            process_file_content(&mut reader, 0, data.len() as u64, Vec::new(), max_buf_size)
2066                .await
2067                .unwrap();
2068
2069        assert_eq!(result.lines.len(), 2);
2070        assert_eq!(result.lines[0], b"line1");
2071        assert_eq!(result.lines[1], b"line2");
2072        assert_eq!(result.new_position, data.len() as u64);
2073        assert_eq!(result.incomplete_line_buffer, b"partial");
2074    }
2075
2076    #[tokio::test]
2077    async fn test_process_file_content_with_existing_buffer() {
2078        let data = b"omplete\nline2\nline3\n".to_vec();
2079        let mut reader = create_mock_reader(data.clone());
2080        let existing_buffer = b"inc".to_vec();
2081        let max_buf_size = 10;
2082
2083        let result = process_file_content(
2084            &mut reader,
2085            0,
2086            data.len() as u64,
2087            existing_buffer,
2088            max_buf_size,
2089        )
2090        .await
2091        .unwrap();
2092
2093        assert_eq!(result.lines.len(), 3);
2094        assert_eq!(result.lines[0], b"incomplete");
2095        assert_eq!(result.lines[1], b"line2");
2096        assert_eq!(result.lines[2], b"line3");
2097        assert_eq!(result.new_position, data.len() as u64);
2098        assert!(result.incomplete_line_buffer.is_empty());
2099    }
2100
2101    #[tokio::test]
2102    async fn test_process_file_content_empty_file() {
2103        let data = Vec::new();
2104        let mut reader = create_mock_reader(data.clone());
2105        let max_buf_size = 10;
2106
2107        let result = process_file_content(&mut reader, 0, 0, Vec::new(), max_buf_size)
2108            .await
2109            .unwrap();
2110
2111        assert!(result.lines.is_empty());
2112        assert_eq!(result.new_position, 0);
2113        assert!(result.incomplete_line_buffer.is_empty());
2114    }
2115
2116    #[tokio::test]
2117    async fn test_process_file_content_only_newlines() {
2118        let data = b"\n\n\n".to_vec();
2119        let mut reader = create_mock_reader(data.clone());
2120        let max_buf_size = 10;
2121
2122        let result =
2123            process_file_content(&mut reader, 0, data.len() as u64, Vec::new(), max_buf_size)
2124                .await
2125                .unwrap();
2126
2127        // Empty lines should not be added (the function skips empty line_buffer)
2128        assert!(result.lines.is_empty());
2129        assert_eq!(result.new_position, data.len() as u64);
2130        assert!(result.incomplete_line_buffer.is_empty());
2131    }
2132
2133    #[tokio::test]
2134    async fn test_process_file_content_no_newlines() {
2135        let data = b"no newlines here".to_vec();
2136        let mut reader = create_mock_reader(data.clone());
2137        let max_buf_size = 10;
2138
2139        let result =
2140            process_file_content(&mut reader, 0, data.len() as u64, Vec::new(), max_buf_size)
2141                .await
2142                .unwrap();
2143
2144        assert!(result.lines.is_empty());
2145        assert_eq!(result.new_position, data.len() as u64);
2146        assert_eq!(result.incomplete_line_buffer, b"no newlines here");
2147    }
2148
2149    #[tokio::test]
2150    async fn test_process_file_content_file_truncation() {
2151        let data = b"line1\nline2\n".to_vec();
2152        let mut reader = create_mock_reader(data.clone());
2153
2154        // Simulate current position being beyond file size (file was truncated)
2155        let result = process_file_content(
2156            &mut reader,
2157            100, // position beyond file size
2158            data.len() as u64,
2159            Vec::new(),
2160            10, // max_buf_size
2161        )
2162        .await
2163        .unwrap();
2164
2165        // Should reset to beginning and read all lines
2166        assert_eq!(result.lines.len(), 2);
2167        assert_eq!(result.lines[0], b"line1");
2168        assert_eq!(result.lines[1], b"line2");
2169        assert_eq!(result.new_position, data.len() as u64);
2170        assert!(result.incomplete_line_buffer.is_empty());
2171    }
2172
2173    #[tokio::test]
2174    async fn test_process_file_content_seek_to_position() {
2175        let data = b"line1\nline2\nline3\n".to_vec();
2176        let mut reader = create_mock_reader(data.clone());
2177
2178        // Start reading from position 6 (after "line1\n")
2179        let result = process_file_content(&mut reader, 6, data.len() as u64, Vec::new(), 10)
2180            .await
2181            .unwrap();
2182
2183        assert_eq!(result.lines.len(), 2);
2184        assert_eq!(result.lines[0], b"line2");
2185        assert_eq!(result.lines[1], b"line3");
2186        assert_eq!(result.new_position, data.len() as u64);
2187        assert!(result.incomplete_line_buffer.is_empty());
2188    }
2189
2190    #[tokio::test]
2191    async fn test_process_file_content_position_equals_file_size() {
2192        let data = b"line1\nline2\n".to_vec();
2193        let mut reader = create_mock_reader(data.clone());
2194
2195        // Start reading from end of file
2196        let result = process_file_content(
2197            &mut reader,
2198            data.len() as u64,
2199            data.len() as u64,
2200            Vec::new(),
2201            10,
2202        )
2203        .await
2204        .unwrap();
2205
2206        // Should not read anything new
2207        assert!(
2208            result.lines.is_empty(),
2209            "Expected empty line got {:?}",
2210            result.lines
2211        );
2212        assert_eq!(result.new_position, data.len() as u64);
2213        assert!(result.incomplete_line_buffer.is_empty());
2214    }
2215
2216    #[tokio::test]
2217    async fn test_process_file_content_large_line_truncation() {
2218        // Create a line longer than MAX_LINE_SIZE
2219        let large_line = "x".repeat(MAX_LINE_SIZE + 1000);
2220        let data = format!("{}\nline2\n", large_line).into_bytes();
2221        let mut reader = create_mock_reader(data.clone());
2222
2223        let result = process_file_content(&mut reader, 0, data.len() as u64, Vec::new(), 10)
2224            .await
2225            .unwrap();
2226
2227        assert_eq!(result.lines.len(), 2);
2228
2229        // First line should be truncated
2230        assert_eq!(
2231            result.lines[0].len(),
2232            MAX_LINE_SIZE + b"... [TRUNCATED]".len()
2233        );
2234        assert!(result.lines[0].ends_with(b"... [TRUNCATED]"));
2235
2236        // Second line should be normal
2237        assert_eq!(result.lines[1], b"line2");
2238
2239        assert_eq!(result.new_position, data.len() as u64);
2240        assert!(result.incomplete_line_buffer.is_empty());
2241    }
2242
2243    #[tokio::test]
2244    async fn test_process_file_content_mixed_line_endings() {
2245        let data = b"line1\nline2\r\nline3\n".to_vec();
2246        let mut reader = create_mock_reader(data.clone());
2247
2248        let result = process_file_content(&mut reader, 0, data.len() as u64, Vec::new(), 10)
2249            .await
2250            .unwrap();
2251
2252        assert_eq!(result.lines.len(), 3);
2253        assert_eq!(result.lines[0], b"line1");
2254        assert_eq!(result.lines[1], b"line2\r"); // \r is preserved
2255        assert_eq!(result.lines[2], b"line3");
2256        assert_eq!(result.new_position, data.len() as u64);
2257        assert!(result.incomplete_line_buffer.is_empty());
2258    }
2259
2260    #[tokio::test]
2261    async fn test_process_file_content_existing_buffer_with_truncation() {
2262        // Create a scenario where existing buffer + new data creates a line that needs truncation
2263        let existing_buffer = "x".repeat(MAX_LINE_SIZE - 100);
2264        let data = format!("{}\nline2\n", "y".repeat(200)).into_bytes();
2265        let mut reader = create_mock_reader(data.clone());
2266
2267        let result = process_file_content(
2268            &mut reader,
2269            0,
2270            data.len() as u64,
2271            existing_buffer.into_bytes(),
2272            10,
2273        )
2274        .await
2275        .unwrap();
2276
2277        assert_eq!(result.lines.len(), 2);
2278
2279        // First line should be truncated (existing buffer + new data)
2280        assert_eq!(
2281            result.lines[0].len(),
2282            MAX_LINE_SIZE + b"... [TRUNCATED]".len()
2283        );
2284        assert!(result.lines[0].ends_with(b"... [TRUNCATED]"));
2285
2286        // Second line should be normal
2287        assert_eq!(result.lines[1], b"line2");
2288
2289        assert_eq!(result.new_position, data.len() as u64);
2290        assert!(result.incomplete_line_buffer.is_empty());
2291    }
2292
2293    #[tokio::test]
2294    async fn test_process_file_content_single_character_lines() {
2295        let data = b"a\nb\nc\n".to_vec();
2296        let mut reader = create_mock_reader(data.clone());
2297
2298        let result = process_file_content(&mut reader, 0, data.len() as u64, Vec::new(), 10)
2299            .await
2300            .unwrap();
2301
2302        assert_eq!(result.lines.len(), 3);
2303        assert_eq!(result.lines[0], b"a");
2304        assert_eq!(result.lines[1], b"b");
2305        assert_eq!(result.lines[2], b"c");
2306        assert_eq!(result.new_position, data.len() as u64);
2307        assert!(result.incomplete_line_buffer.is_empty());
2308    }
2309
2310    #[tokio::test]
2311    async fn test_process_file_content_binary_data() {
2312        let data = vec![0x00, 0x01, 0x02, b'\n', 0xFF, 0xFE, b'\n'];
2313        let mut reader = create_mock_reader(data.clone());
2314
2315        let result = process_file_content(&mut reader, 0, data.len() as u64, Vec::new(), 10)
2316            .await
2317            .unwrap();
2318
2319        assert_eq!(result.lines.len(), 2);
2320        assert_eq!(result.lines[0], vec![0x00, 0x01, 0x02]);
2321        assert_eq!(result.lines[1], vec![0xFF, 0xFE]);
2322        assert_eq!(result.new_position, data.len() as u64);
2323        assert!(result.incomplete_line_buffer.is_empty());
2324    }
2325
2326    #[tokio::test]
2327    async fn test_process_file_content_resume_after_max_buffer_size() {
2328        // Test data: 3 lines as specified in the example
2329        let data = b"line 1\nline 2\nline 3\n".to_vec();
2330        let mut reader = create_mock_reader(data.clone());
2331        let max_buffer_size = 2; // Limit to 2 lines per call
2332
2333        // First call: should return first 2 lines
2334        let result1 = process_file_content(
2335            &mut reader,
2336            0, // start from beginning
2337            data.len() as u64,
2338            Vec::new(), // no existing buffer
2339            max_buffer_size,
2340        )
2341        .await
2342        .unwrap();
2343
2344        // Verify first call results
2345        assert_eq!(result1.lines.len(), 2, "First call should return 2 lines");
2346        assert_eq!(result1.lines[0], b"line 1");
2347        assert_eq!(result1.lines[1], b"line 2");
2348        assert!(result1.incomplete_line_buffer.is_empty());
2349
2350        // The position should be after "line 1\nline 2\n" (14 bytes)
2351        let expected_position_after_first_call = b"line 1\nline 2\n".len() as u64;
2352        assert_eq!(result1.new_position, expected_position_after_first_call);
2353
2354        // Second call: resume from where first call left off
2355        let mut reader2 = create_mock_reader(data.clone());
2356        let result2 = process_file_content(
2357            &mut reader2,
2358            result1.new_position, // resume from previous position
2359            data.len() as u64,
2360            result1.incomplete_line_buffer, // pass any incomplete buffer (should be empty)
2361            max_buffer_size,
2362        )
2363        .await
2364        .unwrap();
2365
2366        // Verify second call results
2367        assert_eq!(result2.lines.len(), 1, "Second call should return 1 line");
2368        assert_eq!(result2.lines[0], b"line 3");
2369        assert!(result2.incomplete_line_buffer.is_empty());
2370        assert_eq!(result2.new_position, data.len() as u64);
2371    }
2372
2373    #[tokio::test]
2374    async fn test_utf8_truncation() {
2375        // Test that StreamFwder doesn't panic when truncating lines
2376        // with multi-byte chars.
2377
2378        hyperactor_telemetry::initialize_logging_for_test();
2379
2380        // Create a line longer than MAX_LINE_SIZE with an emoji at the boundary
2381        let mut long_line = "x".repeat(MAX_LINE_SIZE - 1);
2382        long_line.push('🦀'); // 4-byte emoji - truncation will land in the middle
2383        long_line.push('\n');
2384
2385        // Create IO streams
2386        let (mut writer, reader) = tokio::io::duplex(8192);
2387
2388        // Start StreamFwder
2389        let test_proc_id = test_proc_id("testproc_0");
2390        let monitor = StreamFwder::start_with_writer(
2391            reader,
2392            Box::new(tokio::io::sink()), // discard output
2393            None,                        // no file monitor needed
2394            OutputTarget::Stdout,
2395            1,    // tail buffer of 1 (need at least one sink)
2396            None, // no log channel
2397            &test_proc_id,
2398            None, // no prefix
2399        );
2400
2401        // Write the problematic line
2402        writer.write_all(long_line.as_bytes()).await.unwrap();
2403        drop(writer); // Close to signal EOF
2404
2405        // Wait for completion - should NOT panic
2406        let (_lines, result) = monitor.abort().await;
2407        result.expect("Should complete without panic despite UTF-8 truncation");
2408    }
2409}