dial9_core/worker/
processors.rs

1//! Built-in segment processors: gzip compression and disk write-back.
2
3use crate::payload::Payload;
4use crate::pipeline::{ProcessError, SegmentData, SegmentProcessor};
5use crate::rate_limit::rate_limited;
6use std::future::Future;
7use std::path::PathBuf;
8use std::pin::Pin;
9use std::time::Duration;
10
11/// Gzips the segment payload in-memory. Sets the `content_encoding` and
12/// `write_back_extension` metadata keys so downstream stages know the
13/// payload is gzipped. Already-gzipped segments (detected by magic bytes)
14/// pass through unchanged.
15#[derive(Debug, Default)]
16pub struct GzipCompressor;
17
18impl SegmentProcessor for GzipCompressor {
19    fn name(&self) -> &'static str {
20        "Gzip"
21    }
22
23    fn process(
24        &mut self,
25        mut data: SegmentData,
26    ) -> Pin<Box<dyn Future<Output = Result<SegmentData, ProcessError>> + Send + '_>> {
27        Box::pin(async move {
28            // Skip already-compressed segments to avoid double-gzip.
29            if data.payload().starts_with(&[0x1f, 0x8b]) {
30                data.metadata_mut()
31                    .insert("content_encoding".into(), "gzip".into());
32                data.metadata_mut()
33                    .insert("write_back_extension".into(), ".gz".into());
34                return Ok(data);
35            }
36            let raw = data.take_payload();
37            let compressed = tokio::task::spawn_blocking(move || {
38                use flate2::write::GzEncoder;
39                use std::io::Write;
40                let mut encoder = GzEncoder::new(Vec::new(), flate2::Compression::fast());
41                for chunk in raw.chunks() {
42                    encoder.write_all(chunk)?;
43                }
44                encoder.finish()
45            })
46            .await;
47            match compressed {
48                Ok(Ok(bytes)) => {
49                    data.set_compressed_size(bytes.len() as u64);
50                    data.set_payload(Payload::from_vec(bytes));
51                    data.metadata_mut()
52                        .insert("content_encoding".into(), "gzip".into());
53                    data.metadata_mut()
54                        .insert("write_back_extension".into(), ".gz".into());
55                    Ok(data)
56                }
57                Ok(Err(e)) => Err(ProcessError::io(data, e)),
58                Err(e) => Err(ProcessError::io(data, std::io::Error::other(e))),
59            }
60        })
61    }
62}
63
64/// Writes the current payload bytes back to disk. If a
65/// `write_back_extension` metadata key is present, the bytes are written to
66/// `{original}{extension}` and the original segment file is removed.
67/// When `dir` is set, the file is written to that directory instead of
68/// alongside the original.
69#[derive(Debug, Default)]
70pub struct WriteBackProcessor {
71    dir: Option<PathBuf>,
72}
73
74impl WriteBackProcessor {
75    /// Write to `dir` instead of alongside the original segment.
76    pub fn to_dir(dir: PathBuf) -> Self {
77        Self { dir: Some(dir) }
78    }
79}
80
81impl SegmentProcessor for WriteBackProcessor {
82    fn name(&self) -> &'static str {
83        "WriteBack"
84    }
85
86    fn process(
87        &mut self,
88        data: SegmentData,
89    ) -> Pin<Box<dyn Future<Output = Result<SegmentData, ProcessError>> + Send + '_>> {
90        let output_dir = self.dir.clone();
91        Box::pin(async move {
92            let original_path = match data.segment().disk_path() {
93                Some(p) => p.to_owned(),
94                None => {
95                    return Err(ProcessError::io(
96                        data,
97                        std::io::Error::other(
98                            "WriteBackProcessor requires a disk-backed segment; \
99                             memory-backed segments must not use write_back()",
100                        ),
101                    ));
102                }
103            };
104            let base_path = match &output_dir {
105                Some(dir) => dir.join(original_path.file_name().unwrap_or_default()),
106                None => original_path.clone(),
107            };
108            let dest_path = match data.metadata().get("write_back_extension") {
109                Some(ext) => {
110                    let mut p = base_path.as_os_str().to_owned();
111                    p.push(ext);
112                    std::path::PathBuf::from(p)
113                }
114                None => base_path,
115            };
116            let payload = data.payload().clone();
117            let write_dest = dest_path.clone();
118            let result = tokio::task::spawn_blocking(move || {
119                use std::io::{BufWriter, Write};
120                if let Some(parent) = write_dest.parent() {
121                    std::fs::create_dir_all(parent)?;
122                }
123                let mut f = BufWriter::new(std::fs::File::create(&write_dest)?);
124                for chunk in payload.chunks() {
125                    f.write_all(chunk)?;
126                }
127                f.flush()
128            })
129            .await;
130            match result {
131                Ok(Ok(())) => {
132                    if dest_path != original_path {
133                        // Remove the original .bin now that the output exists elsewhere.
134                        // If the writer already evicted it, clean up the dest
135                        // file we just wrote so it doesn't leak on disk.
136                        match std::fs::remove_file(&original_path) {
137                            Ok(()) => {}
138                            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
139                                let _ = std::fs::remove_file(&dest_path);
140                            }
141                            Err(e) => {
142                                rate_limited!(Duration::from_secs(60), {
143                                    tracing::warn!(
144                                        "failed to remove original segment {}: {e}",
145                                        original_path.display()
146                                    );
147                                });
148                            }
149                        }
150                    }
151                    Ok(data)
152                }
153                Ok(Err(e)) => Err(ProcessError::io(data, e)),
154                Err(e) => Err(ProcessError::io(data, std::io::Error::other(e))),
155            }
156        })
157    }
158}
dial9_core/worker/processors.rs

dial9_core/worker/
processors.rs
