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

rskit_dataset/
sink.rs
