use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use chrono::{DateTime, NaiveDate, TimeZone, Utc};
use tokio::sync::Mutex;
use web_sys::FileSystemDirectoryHandle;
use crate::opfs_helpers::{
opfs_append_quota_guard, opfs_read_all, opfs_root, opfs_walk_or_create_stream, OpfsError,
};
#[derive(Debug, Clone)]
pub struct StorageConfig {
pub root: PathBuf,
pub default_retention_days: u32,
pub orderbook_snapshot_interval_secs: u64,
}
impl Default for StorageConfig {
fn default() -> Self {
Self {
root: PathBuf::from("dig3_storage"),
default_retention_days: 30,
orderbook_snapshot_interval_secs: 300,
}
}
}
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct StreamKey {
pub exchange: String,
pub account: String,
pub symbol: String,
pub stream_kind: String,
}
struct WasmStreamWriter {
sym_dir: FileSystemDirectoryHandle,
day_bufs: HashMap<String, Vec<u8>>,
#[allow(dead_code)]
key: StreamKey,
}
impl WasmStreamWriter {
fn push(&mut self, ts_ms: i64, payload: &[u8]) {
let day = record_day(ts_ms);
let buf = self.day_bufs.entry(day).or_default();
buf.extend_from_slice(&ts_ms.to_le_bytes());
buf.extend_from_slice(&(payload.len() as u32).to_le_bytes());
buf.extend_from_slice(payload);
}
async fn flush(&mut self) -> Result<(), OpfsError> {
for (day, bytes) in self.day_bufs.drain() {
if bytes.is_empty() {
continue;
}
let fname = format!("{day}.bin");
opfs_append_quota_guard(&self.sym_dir, &fname, &bytes).await?;
}
Ok(())
}
}
pub struct StorageManager {
config: StorageConfig,
open_writers: Arc<Mutex<HashMap<StreamKey, WasmStreamWriter>>>,
}
impl StorageManager {
pub fn new(config: StorageConfig) -> std::io::Result<Self> {
Ok(Self {
config,
open_writers: Arc::new(Mutex::new(HashMap::new())),
})
}
pub async fn append(
&self,
key: &StreamKey,
ts_ms: i64,
payload: &[u8],
) -> std::io::Result<()> {
let mut writers = self.open_writers.lock().await;
if !writers.contains_key(key) {
let writer = open_writer(key).await.map_err(std::io::Error::from)?;
writers.insert(key.clone(), writer);
}
writers
.get_mut(key)
.expect("just inserted")
.push(ts_ms, payload);
Ok(())
}
pub async fn flush_all(&self) -> std::io::Result<()> {
let mut writers = self.open_writers.lock().await;
for w in writers.values_mut() {
w.flush().await.map_err(std::io::Error::from)?;
}
Ok(())
}
pub async fn read_range(
&self,
key: &StreamKey,
from_ms: i64,
to_ms: i64,
) -> std::io::Result<Vec<(i64, Vec<u8>)>> {
let from_ms = from_ms.max(0);
let to_ms = to_ms.clamp(0, MAX_SAFE_MS);
if from_ms > to_ms {
return Ok(vec![]);
}
let from_day =
ms_to_date(from_ms).ok_or_else(|| std::io::Error::other("bad from_ms timestamp"))?;
let to_day =
ms_to_date(to_ms).ok_or_else(|| std::io::Error::other("bad to_ms timestamp"))?;
let sym_dir = {
let mut writers = self.open_writers.lock().await;
if !writers.contains_key(key) {
let writer = open_writer(key).await.map_err(std::io::Error::from)?;
writers.insert(key.clone(), writer);
}
writers
.get(key)
.expect("just inserted")
.sym_dir
.clone()
};
let mut out: Vec<(i64, Vec<u8>)> = Vec::new();
let mut day = from_day;
while day <= to_day {
let fname = format!("{}.bin", day.format("%Y-%m-%d"));
match opfs_read_all(&sym_dir, &fname).await {
Ok(bytes) => {
parse_records(&bytes, from_ms, to_ms, &mut out);
}
Err(OpfsError::FileNotFound) => {
}
Err(e) => {
return Err(std::io::Error::from(e));
}
}
day = day.succ_opt().ok_or_else(|| {
std::io::Error::other("date overflow iterating day range")
})?;
}
Ok(out)
}
pub fn stream_dir(&self, key: &StreamKey) -> PathBuf {
self.config
.root
.join(&key.exchange)
.join(&key.account)
.join(&key.symbol)
.join(&key.stream_kind)
}
pub fn cleanup(&self, _now: DateTime<Utc>) -> std::io::Result<usize> {
Ok(0)
}
pub fn config(&self) -> &StorageConfig {
&self.config
}
}
async fn open_writer(key: &StreamKey) -> Result<WasmStreamWriter, OpfsError> {
let root = opfs_root().await?;
let sym_dir = open_stream_dir(&root, key).await?;
Ok(WasmStreamWriter {
sym_dir,
day_bufs: HashMap::new(),
key: key.clone(),
})
}
async fn open_stream_dir(
root: &FileSystemDirectoryHandle,
key: &StreamKey,
) -> Result<FileSystemDirectoryHandle, OpfsError> {
opfs_walk_or_create_stream(
root,
&key.exchange,
&key.account,
&key.symbol,
&key.stream_kind,
)
.await
}
fn parse_records(bytes: &[u8], from_ms: i64, to_ms: i64, out: &mut Vec<(i64, Vec<u8>)>) {
let mut pos = 0usize;
while pos + 12 <= bytes.len() {
let ts = i64::from_le_bytes(
bytes[pos..pos + 8]
.try_into()
.unwrap_or([0u8; 8]),
);
let len = u32::from_le_bytes(
bytes[pos + 8..pos + 12]
.try_into()
.unwrap_or([0u8; 4]),
) as usize;
pos += 12;
if pos + len > bytes.len() {
break; }
let payload = bytes[pos..pos + len].to_vec();
pos += len;
if ts >= from_ms && ts <= to_ms {
out.push((ts, payload));
}
}
}
pub const MAX_SAFE_MS: i64 = 253_402_300_799_999;
fn ms_to_date(ts_ms: i64) -> Option<NaiveDate> {
Utc.timestamp_millis_opt(ts_ms)
.single()
.map(|dt| dt.date_naive())
}
fn record_day(ts_ms: i64) -> String {
ms_to_date(ts_ms.clamp(0, MAX_SAFE_MS))
.unwrap_or_else(|| NaiveDate::from_ymd_opt(1970, 1, 1).expect("epoch is valid"))
.format("%Y-%m-%d")
.to_string()
}