tatara_engine/client/
log_collector.rs

1use anyhow::{Context, Result};
2use chrono::Utc;
3use std::path::PathBuf;
4use tokio::sync::mpsc;
5use tracing::debug;
6
7use crate::drivers::LogEntry;
8
9pub struct LogCollector {
10    alloc_dir: PathBuf,
11}
12
13impl LogCollector {
14    pub fn new(alloc_dir: PathBuf) -> Self {
15        Self { alloc_dir }
16    }
17
18    /// Read existing logs for an allocation's task.
19    pub async fn read_logs(&self, alloc_id: &str, task_name: &str) -> Result<Vec<LogEntry>> {
20        let task_dir = self.alloc_dir.join(alloc_id).join(task_name);
21        let mut entries = Vec::new();
22
23        for (stream, filename) in [("stdout", "stdout.log"), ("stderr", "stderr.log")] {
24            let path = task_dir.join(filename);
25            if path.exists() {
26                let content = tokio::fs::read_to_string(&path)
27                    .await
28                    .with_context(|| format!("Failed to read {}", path.display()))?;
29
30                for line in content.lines() {
31                    entries.push(LogEntry {
32                        task_name: task_name.to_string(),
33                        message: line.to_string(),
34                        stream: stream.to_string(),
35                        timestamp: Utc::now(),
36                    });
37                }
38            }
39        }
40
41        Ok(entries)
42    }
43
44    /// Stream logs by tailing log files. Returns a channel of log entries.
45    pub async fn tail_logs(
46        &self,
47        alloc_id: &str,
48        task_name: &str,
49    ) -> Result<mpsc::Receiver<LogEntry>> {
50        let (tx, rx) = mpsc::channel(256);
51        let stdout_path = self
52            .alloc_dir
53            .join(alloc_id)
54            .join(task_name)
55            .join("stdout.log");
56        let stderr_path = self
57            .alloc_dir
58            .join(alloc_id)
59            .join(task_name)
60            .join("stderr.log");
61        let task_name_stdout = task_name.to_string();
62        let task_name_stderr = task_name.to_string();
63        let tx_stderr = tx.clone();
64
65        tokio::spawn(async move {
66            tail_file(stdout_path, "stdout", &task_name_stdout, tx).await;
67        });
68
69        tokio::spawn(async move {
70            tail_file(stderr_path, "stderr", &task_name_stderr, tx_stderr).await;
71        });
72
73        Ok(rx)
74    }
75}
76
77async fn tail_file(path: PathBuf, stream: &str, task_name: &str, tx: mpsc::Sender<LogEntry>) {
78    use tokio::io::{AsyncBufReadExt, BufReader};
79
80    // Wait for file to exist
81    loop {
82        if path.exists() {
83            break;
84        }
85        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
86    }
87
88    let file = match tokio::fs::File::open(&path).await {
89        Ok(f) => f,
90        Err(e) => {
91            debug!(path = %path.display(), error = %e, "Failed to open log file");
92            return;
93        }
94    };
95
96    let mut reader = BufReader::new(file).lines();
97    let stream = stream.to_string();
98    let task_name = task_name.to_string();
99
100    loop {
101        match reader.next_line().await {
102            Ok(Some(line)) => {
103                let entry = LogEntry {
104                    task_name: task_name.clone(),
105                    message: line,
106                    stream: stream.clone(),
107                    timestamp: Utc::now(),
108                };
109                if tx.send(entry).await.is_err() {
110                    return;
111                }
112            }
113            Ok(None) => {
114                // EOF — wait and retry (tail -f behavior)
115                tokio::time::sleep(std::time::Duration::from_millis(250)).await;
116            }
117            Err(_) => return,
118        }
119    }
120}
tatara_engine/client/log_collector.rs

tatara_engine/client/
log_collector.rs
