digdigdig3_station/storage/
manager.rs

1//! StorageManager — top-level storage owner.
2//!
3//! Manages multiple per-`(exchange, account, symbol, stream_kind)` log streams.
4//! Handles file lifecycle, rotation, and retention.
5//!
6//! All write operations are async-safe: each `RotatingWriter` is wrapped in
7//! `tokio::sync::Mutex` and accessed through an `Arc`.
8
9use std::collections::HashMap;
10use std::path::PathBuf;
11use std::sync::Arc;
12
13use chrono::{DateTime, NaiveDate, TimeZone, Utc};
14use tokio::sync::Mutex;
15
16use super::retention;
17use super::rotation::{read_file_range, RotatingWriter};
18
19// ── Config ────────────────────────────────────────────────────────────────────
20
21/// Top-level configuration for `StorageManager`.
22#[derive(Debug, Clone)]
23pub struct StorageConfig {
24    /// Root directory. Layout: `{root}/{exchange}/{account}/{symbol}/{stream_kind}/`
25    pub root: PathBuf,
26    /// How many days to keep daily files before auto-deletion. Default: 30.
27    pub default_retention_days: u32,
28    /// How often (seconds) to write orderbook snapshots. Default: 300 (5 min).
29    pub orderbook_snapshot_interval_secs: u64,
30}
31
32impl Default for StorageConfig {
33    fn default() -> Self {
34        Self {
35            root: PathBuf::from("dig3_storage"),
36            default_retention_days: 30,
37            orderbook_snapshot_interval_secs: 300,
38        }
39    }
40}
41
42// ── StreamKey ─────────────────────────────────────────────────────────────────
43
44/// Unique identifier for one data stream.
45#[derive(Debug, Clone, Eq, PartialEq, Hash)]
46pub struct StreamKey {
47    pub exchange: String,
48    pub account: String,
49    pub symbol: String,
50    pub stream_kind: String,
51}
52
53// ── StorageManager ────────────────────────────────────────────────────────────
54
55/// Manages rotating daily log files for multiple concurrent streams.
56///
57/// # Example
58/// ```no_run
59/// use digdigdig3::core::storage::{StorageManager, StorageConfig, StreamKey};
60///
61/// #[tokio::main]
62/// async fn main() -> std::io::Result<()> {
63///     let mgr = StorageManager::new(StorageConfig::default())?;
64///     let key = StreamKey {
65///         exchange: "binance".into(),
66///         account: "spot".into(),
67///         symbol: "BTCUSDT".into(),
68///         stream_kind: "ticker".into(),
69///     };
70///     mgr.append(&key, 1_700_000_000_000, b"payload").await?;
71///     Ok(())
72/// }
73/// ```
74pub struct StorageManager {
75    config: StorageConfig,
76    open_writers: Arc<Mutex<HashMap<StreamKey, Arc<Mutex<RotatingWriter>>>>>,
77}
78
79impl StorageManager {
80    /// Create a `StorageManager`. Creates `config.root` if it doesn't exist.
81    pub fn new(config: StorageConfig) -> std::io::Result<Self> {
82        std::fs::create_dir_all(&config.root)?;
83        Ok(Self {
84            config,
85            open_writers: Arc::new(Mutex::new(HashMap::new())),
86        })
87    }
88
89    /// Append one record to the stream identified by `key`.
90    ///
91    /// Rotates to a new daily file automatically at UTC midnight.
92    pub async fn append(&self, key: &StreamKey, ts_ms: i64, payload: &[u8]) -> std::io::Result<()> {
93        let writer = self.get_or_open(key).await?;
94        let mut w = writer.lock().await;
95        w.append(ts_ms, payload)
96    }
97
98    /// Flush all open writers to the OS.
99    pub async fn flush_all(&self) -> std::io::Result<()> {
100        let writers = self.open_writers.lock().await;
101        for w in writers.values() {
102            w.lock().await.flush()?;
103        }
104        Ok(())
105    }
106
107    /// Read records in `[from_ms, to_ms]` (inclusive) for `key`.
108    ///
109    /// Spans multiple daily files automatically. Accepts any `i64` (including
110    /// `i64::MAX` sentinel for "read to end") — values outside the
111    /// representable date range are clamped to `[0, MAX_SAFE_MS]`.
112    pub async fn read_range(
113        &self,
114        key: &StreamKey,
115        from_ms: i64,
116        to_ms: i64,
117    ) -> std::io::Result<Vec<(i64, Vec<u8>)>> {
118        let dir = self.stream_dir(key);
119        if !dir.exists() {
120            return Ok(vec![]);
121        }
122
123        let from_ms = from_ms.max(0);
124        let to_ms = to_ms.clamp(0, MAX_SAFE_MS);
125        if from_ms > to_ms {
126            return Ok(vec![]);
127        }
128
129        let from_day = ms_to_date(from_ms)
130            .ok_or_else(|| std::io::Error::other("bad from_ms timestamp"))?;
131        let to_day = ms_to_date(to_ms)
132            .ok_or_else(|| std::io::Error::other("bad to_ms timestamp"))?;
133
134        let mut files: Vec<(NaiveDate, PathBuf)> = Vec::new();
135        for entry in std::fs::read_dir(&dir)? {
136            let entry = entry?;
137            let path = entry.path();
138            let Some(stem) = path.file_stem().and_then(|s| s.to_str()) else { continue };
139            if path.extension().and_then(|s| s.to_str()) != Some("bin") { continue }
140            let Ok(date) = NaiveDate::parse_from_str(stem, "%Y-%m-%d") else { continue };
141            if date >= from_day && date <= to_day {
142                files.push((date, path));
143            }
144        }
145        files.sort_by_key(|(d, _)| *d);
146
147        let mut out = Vec::new();
148        for (_, file) in files {
149            let records = read_file_range(&file, from_ms, to_ms)?;
150            out.extend(records);
151        }
152        Ok(out)
153    }
154
155    /// Run retention sweep — delete daily files older than `config.default_retention_days`.
156    ///
157    /// Returns the count of deleted files.
158    pub fn cleanup(&self, now: DateTime<Utc>) -> std::io::Result<usize> {
159        retention::sweep(&self.config.root, now, self.config.default_retention_days)
160    }
161
162    /// Return the directory for a stream: `{root}/{exchange}/{account}/{symbol}/{stream_kind}/`
163    pub fn stream_dir(&self, key: &StreamKey) -> PathBuf {
164        self.config
165            .root
166            .join(&key.exchange)
167            .join(&key.account)
168            .join(&key.symbol)
169            .join(&key.stream_kind)
170    }
171
172    /// Return reference to config.
173    pub fn config(&self) -> &StorageConfig {
174        &self.config
175    }
176
177    // ── internal ──────────────────────────────────────────────────────────────
178
179    async fn get_or_open(
180        &self,
181        key: &StreamKey,
182    ) -> std::io::Result<Arc<Mutex<RotatingWriter>>> {
183        let mut writers = self.open_writers.lock().await;
184        if let Some(w) = writers.get(key) {
185            return Ok(w.clone());
186        }
187        let dir = self.stream_dir(key);
188        let w = Arc::new(Mutex::new(RotatingWriter::new(dir)?));
189        writers.insert(key.clone(), w.clone());
190        Ok(w)
191    }
192}
193
194// ── helpers ───────────────────────────────────────────────────────────────────
195
196/// Largest UTC millisecond timestamp that `chrono` can convert back to a
197/// `NaiveDate` — corresponds to `9999-12-31T23:59:59.999Z`. Anything above
198/// this overflows `ms_to_date`.
199pub const MAX_SAFE_MS: i64 = 253_402_300_799_999;
200
201fn ms_to_date(ts_ms: i64) -> Option<NaiveDate> {
202    Utc.timestamp_millis_opt(ts_ms)
203        .single()
204        .map(|dt| dt.date_naive())
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210
211    struct TmpDir(PathBuf);
212    impl TmpDir {
213        fn new(tag: &str) -> Self {
214            let p = std::env::temp_dir().join(format!(
215                "dig3-mgr-test-{}-{}",
216                tag,
217                Utc::now().timestamp_nanos_opt().unwrap_or(0)
218            ));
219            std::fs::create_dir_all(&p).unwrap();
220            Self(p)
221        }
222        fn path(&self) -> &std::path::Path { &self.0 }
223    }
224    impl Drop for TmpDir {
225        fn drop(&mut self) { let _ = std::fs::remove_dir_all(&self.0); }
226    }
227
228    fn make_key() -> StreamKey {
229        StreamKey {
230            exchange: "binance".into(),
231            account: "spot".into(),
232            symbol: "BTCUSDT".into(),
233            stream_kind: "trade".into(),
234        }
235    }
236
237    #[tokio::test]
238    async fn read_range_clamps_i64_max_sentinel() {
239        let tmp = TmpDir::new("max");
240        let mgr = StorageManager::new(StorageConfig {
241            root: tmp.path().to_path_buf(),
242            ..Default::default()
243        }).unwrap();
244        let key = make_key();
245
246        let now = Utc::now().timestamp_millis();
247        mgr.append(&key, now, b"a").await.unwrap();
248        mgr.append(&key, now + 1, b"b").await.unwrap();
249        mgr.flush_all().await.unwrap();
250
251        // i64::MAX must NOT hang or error — must clamp + return all records.
252        let out = mgr.read_range(&key, 0, i64::MAX).await.unwrap();
253        assert_eq!(out.len(), 2, "expected 2 records back, got {}", out.len());
254        assert_eq!(out[0].1, b"a");
255        assert_eq!(out[1].1, b"b");
256    }
257
258    #[tokio::test]
259    async fn read_range_negative_from_is_clamped() {
260        let tmp = TmpDir::new("neg");
261        let mgr = StorageManager::new(StorageConfig {
262            root: tmp.path().to_path_buf(),
263            ..Default::default()
264        }).unwrap();
265        let key = make_key();
266        let now = Utc::now().timestamp_millis();
267        mgr.append(&key, now, b"r").await.unwrap();
268        mgr.flush_all().await.unwrap();
269
270        let out = mgr.read_range(&key, i64::MIN, i64::MAX).await.unwrap();
271        assert_eq!(out.len(), 1);
272    }
273
274    #[tokio::test]
275    async fn read_range_missing_dir_returns_empty() {
276        let tmp = TmpDir::new("ghost");
277        let mgr = StorageManager::new(StorageConfig {
278            root: tmp.path().to_path_buf(),
279            ..Default::default()
280        }).unwrap();
281        let key = StreamKey {
282            exchange: "ghost".into(), account: "x".into(),
283            symbol: "x".into(), stream_kind: "x".into(),
284        };
285        let out = mgr.read_range(&key, 0, i64::MAX).await.unwrap();
286        assert!(out.is_empty());
287    }
288}
digdigdig3_station/storage/manager.rs

digdigdig3_station/storage/
manager.rs
