rskit_dataset/record/
sink.rs

1//! Record sink adapter — materialize records through a [`DatasetWriter`] as an [`ItemSink`].
2
3use std::path::{Path, PathBuf};
4
5use rskit_errors::{AppError, AppResult, ErrorCode};
6use std::sync::Arc;
7use tokio::sync::{Mutex, mpsc};
8use tokio::task::JoinHandle;
9
10use crate::sink::ItemSink;
11
12use super::model::DatasetRecord;
13use super::ops::BoxRecordStream;
14use super::writer::DatasetWriter;
15
16/// Bounded buffer between the engine's workers and the background writer task.
17const RECORD_CHANNEL_CAPACITY: usize = 128;
18
19/// Adapts any [`DatasetWriter`] into an [`ItemSink`] that streams records to one dataset file.
20///
21/// Records are forwarded over a bounded channel to a background writer task started in
22/// [`prepare`](ItemSink::prepare), preserving backpressure; [`finish`](ItemSink::finish) closes the
23/// channel and awaits the writer.
24pub struct RecordSink {
25    writer: Arc<dyn DatasetWriter>,
26    path: PathBuf,
27    state: Mutex<Option<SinkState>>,
28}
29
30struct SinkState {
31    tx: mpsc::Sender<AppResult<DatasetRecord>>,
32    handle: JoinHandle<AppResult<usize>>,
33}
34
35impl RecordSink {
36    /// Create a sink that writes records to `path` using `writer`.
37    #[must_use]
38    pub fn new(writer: Arc<dyn DatasetWriter>, path: impl Into<PathBuf>) -> Self {
39        Self {
40            writer,
41            path: path.into(),
42            state: Mutex::new(None),
43        }
44    }
45}
46
47#[async_trait::async_trait]
48impl ItemSink<DatasetRecord> for RecordSink {
49    fn name(&self) -> &str {
50        "record"
51    }
52
53    async fn prepare(&self) -> AppResult<()> {
54        if let Some(parent) = self.path.parent()
55            && !parent.as_os_str().is_empty()
56        {
57            let parent = parent.to_path_buf();
58            tokio::task::spawn_blocking(move || {
59                std::fs::create_dir_all(&parent).map_err(|error| {
60                    AppError::new(
61                        ErrorCode::Internal,
62                        format!(
63                            "failed to create dataset directory {}: {error}",
64                            parent.display()
65                        ),
66                    )
67                })
68            })
69            .await
70            .map_err(AppError::internal)??;
71        }
72
73        let (tx, mut rx) = mpsc::channel::<AppResult<DatasetRecord>>(RECORD_CHANNEL_CAPACITY);
74        let stream: BoxRecordStream =
75            Box::pin(futures::stream::poll_fn(move |cx| rx.poll_recv(cx)));
76        let writer = self.writer.clone();
77        let path = self.path.clone();
78        let handle = tokio::spawn(async move { writer.write(stream, &path).await });
79        *self.state.lock().await = Some(SinkState { tx, handle });
80        Ok(())
81    }
82
83    async fn write(&self, item: DatasetRecord) -> AppResult<()> {
84        let tx = {
85            let state = self.state.lock().await;
86            state.as_ref().map(|state| state.tx.clone())
87        };
88        match tx {
89            Some(tx) => tx.send(Ok(item)).await.map_err(|_| {
90                AppError::new(ErrorCode::Internal, "record writer task stopped early")
91            }),
92            None => Err(AppError::new(
93                ErrorCode::Internal,
94                "record sink used before prepare",
95            )),
96        }
97    }
98
99    async fn finish(&self) -> AppResult<()> {
100        let state = self.state.lock().await.take();
101        if let Some(state) = state {
102            drop(state.tx);
103            state.handle.await.map_err(AppError::internal)??;
104        }
105        Ok(())
106    }
107}
108
109impl RecordSink {
110    /// Path this sink writes to.
111    #[must_use]
112    pub fn path(&self) -> &Path {
113        &self.path
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use serde_json::json;
120
121    use super::*;
122    use crate::record::JsonLinesWriter;
123
124    #[tokio::test]
125    async fn record_sink_streams_records_to_file() {
126        let dir = tempfile::tempdir().unwrap();
127        let path = dir.path().join("nested/out.jsonl");
128        let sink = RecordSink::new(Arc::new(JsonLinesWriter), &path);
129
130        sink.prepare().await.unwrap();
131        sink.write(DatasetRecord::from_fields([("id", json!(1))]))
132            .await
133            .unwrap();
134        sink.write(DatasetRecord::from_fields([("id", json!(2))]))
135            .await
136            .unwrap();
137        sink.finish().await.unwrap();
138
139        let written = std::fs::read_to_string(&path).unwrap();
140        assert_eq!(written.lines().count(), 2);
141        assert!(written.contains("\"id\":1"));
142        assert!(written.contains("\"id\":2"));
143    }
144
145    #[tokio::test]
146    async fn record_sink_write_before_prepare_is_rejected() {
147        let dir = tempfile::tempdir().unwrap();
148        let sink = RecordSink::new(Arc::new(JsonLinesWriter), dir.path().join("out.jsonl"));
149        let error = sink
150            .write(DatasetRecord::from_fields([("id", json!(1))]))
151            .await
152            .unwrap_err();
153        assert_eq!(error.code(), ErrorCode::Internal);
154    }
155}
rskit_dataset/record/sink.rs

rskit_dataset/record/
sink.rs
