rskit_dataset/
sink.rs

1//! Per-item materialization seam — where the generic engine hands each collected item to storage.
2//!
3//! The engine counts and routes items via [`DatasetItem`], but never writes them itself.
4//! A [`ItemSink`] owns all item-specific materialization:
5//! [`LocalBlobSink`] writes [`DataItem`] samples to `real/`
6//! and `ai/` directories with offset filenames,
7//! exactly as the collector did before the engine was generalized.
8
9use std::path::{Path, PathBuf};
10use std::sync::atomic::{AtomicUsize, Ordering};
11
12use rskit_errors::{AppError, AppResult, ErrorCode};
13
14use crate::{DataItem, DatasetItem, DatasetLimits, Label};
15
16/// Destination that materializes each collected item as the engine streams it.
17#[async_trait::async_trait]
18pub trait ItemSink<T: DatasetItem>: Send + Sync {
19    /// Stable sink name.
20    fn name(&self) -> &str;
21
22    /// Prepare the destination before any item is written (create directories, seed counters).
23    async fn prepare(&self) -> AppResult<()> {
24        Ok(())
25    }
26
27    /// Materialize a single collected item.
28    async fn write(&self, item: T) -> AppResult<()>;
29
30    /// Flush and finalize the destination after the last item is written.
31    async fn finish(&self) -> AppResult<()> {
32        Ok(())
33    }
34}
35
36/// Local filesystem sink for [`DataItem`] samples.
37///
38/// Writes real samples to `<output>/real` and AI-generated samples to `<output>/ai`,
39/// naming files by a shared,
40/// monotonically increasing counter that resumes past any files already on disk.
41pub struct LocalBlobSink {
42    real_dir: PathBuf,
43    ai_dir: PathBuf,
44    limits: DatasetLimits,
45    file_counter: AtomicUsize,
46}
47
48impl LocalBlobSink {
49    /// Create a sink rooted at `output_dir`, writing into its `real/` and `ai/` subdirectories.
50    #[must_use]
51    pub fn new(output_dir: &Path, limits: DatasetLimits) -> Self {
52        Self {
53            real_dir: output_dir.join("real"),
54            ai_dir: output_dir.join("ai"),
55            limits,
56            file_counter: AtomicUsize::new(0),
57        }
58    }
59}
60
61#[async_trait::async_trait]
62impl ItemSink<DataItem> for LocalBlobSink {
63    fn name(&self) -> &str {
64        "local-blob"
65    }
66
67    async fn prepare(&self) -> AppResult<()> {
68        let real_dir = self.real_dir.clone();
69        let ai_dir = self.ai_dir.clone();
70        let existing = tokio::task::spawn_blocking(move || -> AppResult<usize> {
71            create_dir(&real_dir)?;
72            create_dir(&ai_dir)?;
73            Ok(count_files(&real_dir)? + count_files(&ai_dir)?)
74        })
75        .await
76        .map_err(AppError::internal)??;
77        self.file_counter.store(existing, Ordering::SeqCst);
78        Ok(())
79    }
80
81    async fn write(&self, item: DataItem) -> AppResult<()> {
82        let subdir = if item.label() == Label::Real {
83            &self.real_dir
84        } else {
85            &self.ai_dir
86        };
87        let file_idx = self.file_counter.fetch_add(1, Ordering::SeqCst);
88        let path = subdir.join(format!("{:06}{}", file_idx, item.extension));
89        let limits = self.limits;
90        tokio::task::spawn_blocking(move || item.write_to_path(&path, &limits))
91            .await
92            .map_err(AppError::internal)??;
93        Ok(())
94    }
95}
96
97fn create_dir(dir: &Path) -> AppResult<()> {
98    std::fs::create_dir_all(dir).map_err(|error| {
99        AppError::new(
100            ErrorCode::Internal,
101            format!(
102                "failed to create dataset directory {}: {error}",
103                dir.display()
104            ),
105        )
106    })
107}
108
109/// Count regular files directly under `dir`, treating a missing directory as empty.
110pub(crate) fn count_files(dir: &Path) -> AppResult<usize> {
111    if !dir.exists() {
112        return Ok(0);
113    }
114    std::fs::read_dir(dir)
115        .map_err(|error| {
116            AppError::new(
117                ErrorCode::Internal,
118                format!(
119                    "failed to read dataset output directory {}: {error}",
120                    dir.display()
121                ),
122            )
123        })
124        .map(|entries| {
125            entries
126                .filter_map(Result::ok)
127                .filter(|entry| entry.path().is_file())
128                .count()
129        })
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use crate::MediaType;
136
137    #[test]
138    fn count_files_ignores_directories_and_reports_read_errors() {
139        let dir = tempfile::tempdir().unwrap();
140        std::fs::write(dir.path().join("file.txt"), b"x").unwrap();
141        std::fs::create_dir(dir.path().join("nested")).unwrap();
142        assert_eq!(count_files(dir.path()).unwrap(), 1);
143        assert_eq!(count_files(&dir.path().join("missing")).unwrap(), 0);
144
145        let file = dir.path().join("not-dir");
146        std::fs::write(&file, b"x").unwrap();
147        let err = count_files(&file).unwrap_err();
148        assert_eq!(err.code(), ErrorCode::Internal);
149        assert!(
150            err.to_string()
151                .contains("failed to read dataset output directory")
152        );
153    }
154
155    #[tokio::test]
156    async fn local_blob_sink_prepares_dirs_and_writes_labeled_files() {
157        let out = tempfile::tempdir().unwrap();
158        let sink = LocalBlobSink::new(out.path(), DatasetLimits::default());
159        sink.prepare().await.unwrap();
160
161        let real = DataItem::new(b"r".to_vec(), Label::Real, MediaType::Text, "s")
162            .unwrap()
163            .with_extension(".txt");
164        let ai = DataItem::new(b"a".to_vec(), Label::AiGenerated, MediaType::Text, "s")
165            .unwrap()
166            .with_extension(".txt");
167        sink.write(real).await.unwrap();
168        sink.write(ai).await.unwrap();
169
170        assert_eq!(
171            std::fs::read(out.path().join("real/000000.txt")).unwrap(),
172            b"r"
173        );
174        assert_eq!(
175            std::fs::read(out.path().join("ai/000001.txt")).unwrap(),
176            b"a"
177        );
178    }
179
180    #[tokio::test]
181    async fn local_blob_sink_prepare_reports_directory_creation_errors() {
182        let out = tempfile::tempdir().unwrap();
183        let blocker = out.path().join("blocker");
184        std::fs::write(&blocker, b"x").unwrap();
185        let sink = LocalBlobSink::new(&blocker, DatasetLimits::default());
186        assert_eq!(
187            sink.prepare().await.unwrap_err().code(),
188            ErrorCode::Internal
189        );
190    }
191
192    #[tokio::test]
193    async fn local_blob_sink_counter_resumes_past_existing_files() {
194        let out = tempfile::tempdir().unwrap();
195        std::fs::create_dir_all(out.path().join("real")).unwrap();
196        std::fs::write(out.path().join("real/000000.txt"), b"old").unwrap();
197
198        let sink = LocalBlobSink::new(out.path(), DatasetLimits::default());
199        sink.prepare().await.unwrap();
200        let item = DataItem::new(b"new".to_vec(), Label::Real, MediaType::Text, "s")
201            .unwrap()
202            .with_extension(".txt");
203        sink.write(item).await.unwrap();
204        assert!(out.path().join("real/000001.txt").exists());
205    }
206}
rskit_dataset/sink.rs

rskit_dataset/
sink.rs
