use crate::store_types::{CountsValue, HourTruncatedCursor, SketchSecretPrefix}; use crate::{ error::StorageError, ConsumerInfo, Cursor, EventBatch, JustCount, NsidCount, NsidPrefix, OrderCollectionsBy, PrefixChild, UFOsRecord, }; use async_trait::async_trait; use jetstream::exports::{Did, Nsid}; use metrics::{describe_histogram, histogram, Unit}; use std::collections::{HashMap, HashSet}; use std::path::Path; use std::time::{Duration, Instant}; use tokio::sync::mpsc::Receiver; use tokio_util::sync::CancellationToken; pub type StorageResult = Result; pub trait StorageWhatever, B: StoreBackground, C> { fn init( path: impl AsRef, endpoint: String, force_endpoint: bool, config: C, ) -> StorageResult<(R, W, Option, SketchSecretPrefix)> where Self: Sized; } #[async_trait] pub trait StoreWriter: Clone + Send + Sync where Self: 'static, { fn background_tasks(&mut self, reroll: bool) -> StorageResult; async fn receive_batches( self, mut batches: Receiver>, ) -> StorageResult<()> { describe_histogram!( "storage_slow_batches", Unit::Microseconds, "batches that took more than 3s to insert" ); describe_histogram!( "storage_batch_insert_time", Unit::Microseconds, "total time to insert one commit batch" ); while let Some(event_batch) = batches.recv().await { let token = CancellationToken::new(); let cancelled = token.clone(); tokio::spawn(async move { let started = Instant::now(); let mut concerned = false; loop { tokio::select! { _ = tokio::time::sleep(Duration::from_secs(3)) => { if !concerned { log::warn!("taking a long time to insert an event batch..."); } concerned = true; } _ = cancelled.cancelled() => { if concerned { log::warn!("finally inserted slow event batch (or failed) after {:?}", started.elapsed()); histogram!("storage_slow_batches").record(started.elapsed().as_micros() as f64); } break } } } }); tokio::task::spawn_blocking({ let mut me = self.clone(); move || { let _guard = token.drop_guard(); let t0 = Instant::now(); let r = me.insert_batch(event_batch); histogram!("storage_batch_insert_time").record(t0.elapsed().as_micros() as f64); r } }) .await??; } Err(StorageError::BatchSenderExited) } fn insert_batch( &mut self, event_batch: EventBatch, ) -> StorageResult<()>; fn step_rollup(&mut self) -> StorageResult<(usize, HashSet)>; fn trim_collection( &mut self, collection: &Nsid, limit: usize, full_scan: bool, ) -> StorageResult<(usize, usize, bool)>; fn delete_account(&mut self, did: &Did) -> StorageResult; } #[async_trait] pub trait StoreBackground: Send + Sync { async fn run(mut self, backfill: bool) -> StorageResult<()>; } #[async_trait] pub trait StoreReader: Send + Sync { fn name(&self) -> String; fn update_metrics(&self) {} async fn get_storage_stats(&self) -> StorageResult; async fn get_consumer_info(&self) -> StorageResult; async fn get_collections( &self, limit: usize, order: OrderCollectionsBy, since: Option, until: Option, ) -> StorageResult<(Vec, Option>)>; async fn get_prefix( &self, prefix: NsidPrefix, limit: usize, order: OrderCollectionsBy, since: Option, until: Option, ) -> StorageResult<(JustCount, Vec, Option>)>; async fn get_timeseries( &self, collections: Vec, since: HourTruncatedCursor, until: Option, step: u64, ) -> StorageResult<(Vec, HashMap>)>; async fn get_collection_counts( &self, collection: &Nsid, since: HourTruncatedCursor, until: Option, ) -> StorageResult; async fn get_records_by_collections( &self, collections: HashSet, limit: usize, expand_each_collection: bool, ) -> StorageResult>; async fn search_collections(&self, terms: Vec) -> StorageResult>; }