Stage 2

Author: tzanko-matev

codetracer-python-recorder/Cargo.lock

@@ -23,6 +23,7 @@ default = ["extension-module"]
 [dependencies]
 pyo3 = { version = "0.25.1" }
 runtime_tracing = "0.14.0"
+base64 = "0.22"
 bitflags = "2.4"
 once_cell = "1.19"
 dashmap = "5.5"

codetracer-python-recorder/Cargo.toml

@@ -1,16 +1,19 @@
-#![allow(dead_code)]
-
 //! Cross-platform IO capture workers for stdout, stderr, and stdin.
 
 use std::thread::JoinHandle;
 use std::time::Instant;
 
+use base64::engine::general_purpose::STANDARD as BASE64_ENGINE;
+use base64::Engine;
 use crossbeam_channel::{bounded, Receiver, Sender};
-
 use recorder_errors::{bug, enverr, ErrorCode, RecorderResult};
+use runtime_tracing::{EventLogKind, RecordEvent, TraceLowLevelEvent, TraceWriter};
+use serde::Serialize;
 
 use crate::errors::Result;
 
+use super::thread_snapshots::{SnapshotEntry, ThreadSnapshotStore};
+use super::trace_writer_host::TraceWriterHost;
 use super::IoDrain;
 
 #[cfg(unix)]
@@ -44,7 +47,6 @@ impl StreamKind {
     }
 }
 
-#[allow(dead_code)]
 #[derive(Debug)]
 pub struct IoChunk {
     pub stream: StreamKind,
@@ -105,12 +107,6 @@ impl IoCapture {
     }
 }
 
-impl IoDrain for IoCapture {
-    fn drain(&mut self, _py: pyo3::Python<'_>) -> RecorderResult<()> {
-        self.shutdown()
-    }
-}
-
 impl Drop for IoCapture {
     fn drop(&mut self) {
         if self.platform.is_some() {
@@ -121,6 +117,177 @@ impl Drop for IoCapture {
     }
 }
 
+#[derive(Clone)]
+struct IoEventSink {
+    writer: TraceWriterHost,
+    snapshots: ThreadSnapshotStore,
+    start: Instant,
+}
+
+impl IoEventSink {
+    fn new(writer: TraceWriterHost, snapshots: ThreadSnapshotStore) -> Self {
+        Self {
+            writer,
+            snapshots,
+            start: Instant::now(),
+        }
+    }
+
+    fn pump(self, receiver: IoChunkReceiver) -> Result<()> {
+        let mut sink = self;
+        while let Ok(chunk) = receiver.recv() {
+            sink.record_chunk(chunk)?;
+        }
+        Ok(())
+    }
+
+    fn record_chunk(&mut self, chunk: IoChunk) -> Result<()> {
+        let event = TraceLowLevelEvent::Event(RecordEvent {
+            kind: map_stream_to_event_kind(chunk.stream),
+            metadata: self.serialize_metadata(&chunk)?,
+            content: BASE64_ENGINE.encode(&chunk.bytes),
+        });
+
+        let mut writer = self.writer.lock()?;
+        writer.add_event(event);
+        Ok(())
+    }
+
+    fn serialize_metadata(&self, chunk: &IoChunk) -> Result<String> {
+        let metadata = IoChunkMetadata {
+            stream: chunk.stream.as_str(),
+            timestamp_ns: self.relative_timestamp_ns(chunk.timestamp),
+            thread_id: format!("{:?}", chunk.producer_thread),
+            byte_len: chunk.bytes.len(),
+            snapshot: self.snapshot_metadata(chunk.producer_thread),
+        };
+
+        serde_json::to_string(&metadata).map_err(|err| {
+            bug!(
+                ErrorCode::Unknown,
+                "failed to encode IO chunk metadata as JSON"
+            )
+            .with_context("error", err.to_string())
+        })
+    }
+
+    fn relative_timestamp_ns(&self, timestamp: Instant) -> u128 {
+        match timestamp.checked_duration_since(self.start) {
+            Some(duration) => duration.as_nanos(),
+            None => 0,
+        }
+    }
+
+    fn snapshot_metadata(&self, thread_id: std::thread::ThreadId) -> Option<IoSnapshotMetadata> {
+        let snapshot = self
+            .snapshots
+            .snapshot_for(thread_id)
+            .or_else(|| self.snapshots.latest());
+        snapshot.map(IoSnapshotMetadata::from)
+    }
+}
+
+#[derive(Serialize)]
+struct IoChunkMetadata<'a> {
+    stream: &'a str,
+    timestamp_ns: u128,
+    thread_id: String,
+    byte_len: usize,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    snapshot: Option<IoSnapshotMetadata>,
+}
+
+#[derive(Serialize)]
+struct IoSnapshotMetadata {
+    path_id: usize,
+    line: i64,
+    frame_id: usize,
+}
+
+impl From<SnapshotEntry> for IoSnapshotMetadata {
+    fn from(entry: SnapshotEntry) -> Self {
+        IoSnapshotMetadata {
+            path_id: entry.path_id.0,
+            line: entry.line.0,
+            frame_id: entry.frame_id,
+        }
+    }
+}
+
+fn map_stream_to_event_kind(stream: StreamKind) -> EventLogKind {
+    match stream {
+        StreamKind::Stdout => EventLogKind::Write,
+        StreamKind::Stderr => EventLogKind::WriteOther,
+        StreamKind::Stdin => EventLogKind::Read,
+    }
+}
+
+pub struct ActiveCapture {
+    capture: IoCapture,
+    sink: Option<WorkerHandle>,
+}
+
+impl ActiveCapture {
+    pub fn start(writer: TraceWriterHost, snapshots: ThreadSnapshotStore) -> Result<Self> {
+        let mut capture = IoCapture::start()?;
+        let receiver = capture.take_receiver().ok_or_else(|| {
+            bug!(
+                ErrorCode::Unknown,
+                "IO capture receiver already taken before sink start"
+            )
+        })?;
+
+        let sink_worker = spawn_reader_thread("sink", move || {
+            IoEventSink::new(writer, snapshots).pump(receiver)
+        })?;
+
+        Ok(Self {
+            capture,
+            sink: Some(sink_worker),
+        })
+    }
+
+    fn join_sink(&mut self) -> Result<()> {
+        if let Some(handle) = self.sink.take() {
+            match handle.join() {
+                Ok(result) => result,
+                Err(panic) => {
+                    let message = if let Some(msg) = panic.downcast_ref::<&'static str>() {
+                        (*msg).to_string()
+                    } else if let Some(msg) = panic.downcast_ref::<String>() {
+                        msg.clone()
+                    } else {
+                        "unknown panic".to_string()
+                    };
+                    Err(bug!(ErrorCode::Unknown, "IO capture sink worker panicked")
+                        .with_context("details", message))
+                }
+            }
+        } else {
+            Ok(())
+        }
+    }
+}
+
+impl IoDrain for ActiveCapture {
+    fn drain(&mut self, _py: pyo3::Python<'_>) -> RecorderResult<()> {
+        let shutdown_result = self.capture.shutdown();
+        let sink_result = self.join_sink();
+        shutdown_result.and(sink_result)
+    }
+}
+
+impl Drop for ActiveCapture {
+    fn drop(&mut self) {
+        if let Err(err) = self.capture.shutdown() {
+            log::error!("failed to shutdown IO capture cleanly: {}", err);
+        }
+        if let Err(err) = self.join_sink() {
+            log::error!("failed to join IO sink worker: {}", err);
+        }
+    }
+}
+
 pub(super) fn spawn_reader_thread<F>(name: &'static str, worker: F) -> Result<WorkerHandle>
 where
     F: FnOnce() -> Result<()> + Send + 'static,

codetracer-python-recorder/src/runtime/io_capture/mod.rs

@@ -1,5 +1,3 @@
-#![allow(dead_code)]
-
 use std::fs::File;
 use std::io::{Read, Write};
 use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};

codetracer-python-recorder/src/runtime/io_capture/unix.rs

@@ -1,5 +1,3 @@
-#![allow(dead_code)]
-
 use std::io::ErrorKind;
 
 use crossbeam_channel::Sender;

codetracer-python-recorder/src/runtime/io_capture/windows.rs

@@ -14,6 +14,7 @@ pub use output_paths::TraceOutputPaths;
 
 use activation::ActivationController;
 use frame_inspector::capture_frame;
+use io_capture::ActiveCapture;
 use logging::log_event;
 use thread_snapshots::{SnapshotEntry, ThreadSnapshotStore};
 use trace_writer_host::{register_step_with_guard, TraceWriterHost};
@@ -263,7 +264,7 @@ impl RuntimeTracer {
         self.encountered_failure = false;
         set_active_trace_id(Some(self.trace_id.clone()));
         self.thread_snapshots.reset();
-        self.io_drain = Box::new(NoopIoDrain::default());
+        self.configure_io_capture()?;
         Ok(())
     }
 
@@ -297,6 +298,13 @@ impl RuntimeTracer {
         self.thread_snapshots.clone()
     }
 
+    fn configure_io_capture(&mut self) -> PyResult<()> {
+        let capture = ActiveCapture::start(self.writer.clone(), self.thread_snapshots.clone())
+            .map_err(ffi::map_recorder_error)?;
+        self.io_drain = Box::new(capture);
+        Ok(())
+    }
+
     fn cleanup_partial_outputs(&self) -> RecorderResult<()> {
         if let Some(outputs) = &self.output_paths {
             for path in [outputs.events(), outputs.metadata(), outputs.paths()] {

codetracer-python-recorder/src/runtime/mod.rs

@@ -8,7 +8,6 @@ use log::warn;
 use runtime_tracing::{Line, PathId};
 
 /// Snapshot of the last recorded step for a thread.
-#[allow(dead_code)]
 #[derive(Clone, Debug)]
 pub struct SnapshotEntry {
     pub path_id: PathId,
@@ -46,14 +45,12 @@ impl ThreadSnapshotStore {
     }
 
     /// Return the latest snapshot for a given thread id.
-    #[allow(dead_code)]
     pub fn snapshot_for(&self, thread_id: ThreadId) -> Option<SnapshotEntry> {
         let store = lock(&self.inner);
         store.get(&thread_id).cloned()
     }
 
     /// Return the most recent snapshot observed across all threads.
-    #[allow(dead_code)]
     pub fn latest(&self) -> Option<SnapshotEntry> {
         let global = lock(&self.last_global);
         global.clone()

codetracer-python-recorder/src/runtime/thread_snapshots.rs

@@ -0,0 +1,85 @@
+"""Integration tests for runtime IO capture."""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+import subprocess
+import sys
+import textwrap
+from collections import defaultdict
+from pathlib import Path
+from typing import DefaultDict, Iterable, Tuple
+
+
+def _load_io_events(trace_file: Path) -> Iterable[Tuple[int, dict[str, object], bytes]]:
+    events = json.loads(trace_file.read_text())
+    for entry in events:
+        payload = entry.get("Event")
+        if not payload:
+            continue
+        metadata = json.loads(payload["metadata"])
+        chunk = base64.b64decode(payload["content"])
+        yield payload["kind"], metadata, chunk
+
+
+def test_io_capture_records_all_streams(tmp_path: Path) -> None:
+    script = textwrap.dedent(
+        """
+        import sys
+        import codetracer_python_recorder as codetracer
+
+        sys.stdout.write("hello stdout\\n")
+        sys.stdout.flush()
+        sys.stderr.write("warning\\n")
+        sys.stderr.flush()
+        data = sys.stdin.readline()
+        sys.stdout.write(f"input={data}")
+        sys.stdout.flush()
+        codetracer.stop()
+        """
+    )
+
+    env = os.environ.copy()
+    env["CODETRACER_TRACE"] = str(tmp_path)
+    env["CODETRACER_FORMAT"] = "json"
+
+    completed = subprocess.run(
+        [sys.executable, "-c", script],
+        input="feed\n",
+        text=True,
+        capture_output=True,
+        env=env,
+        check=True,
+    )
+
+    assert completed.stdout == "hello stdout\ninput=feed\n"
+    assert completed.stderr == "warning\n"
+
+    trace_file = tmp_path / "trace.json"
+    assert trace_file.exists(), "expected trace artefact"
+
+    buffers: DefaultDict[str, bytearray] = defaultdict(bytearray)
+
+    for kind, metadata, chunk in _load_io_events(trace_file):
+        stream = metadata["stream"]
+        buffers[stream].extend(chunk)
+
+        assert metadata["byte_len"] == len(chunk)
+        assert metadata["timestamp_ns"] >= 0
+        assert isinstance(metadata["thread_id"], str)
+        snapshot = metadata.get("snapshot")
+        if snapshot is not None:
+            assert {"path_id", "line", "frame_id"}.issubset(snapshot)
+
+        if stream == "stdout":
+            assert kind == 0  # EventLogKind::Write
+        elif stream == "stderr":
+            assert kind == 2  # EventLogKind::WriteOther
+        elif stream == "stdin":
+            assert kind == 3  # EventLogKind::Read
+
+    assert buffers["stdout"].decode() == "hello stdout\ninput=feed\n"
+    assert buffers["stderr"].decode() == "warning\n"
+    assert buffers["stdin"].decode() == "feed\n"