diff --git a/ufos/src/storage_fjall.rs b/ufos/src/storage_fjall.rs index ff5f653..3e0c464 100644 --- a/ufos/src/storage_fjall.rs +++ b/ufos/src/storage_fjall.rs @@ -70,19 +70,44 @@ const MAX_BATCHED_ROLLUP_COUNTS: usize = 256; /// - key: "live_counts" || u64 || nullstr (js_cursor, nsid) /// - val: u64 || HLL (count (not cursor), estimator) /// +/// /// - Hourly total record counts and dids estimate per collection /// - key: "hourly_counts" || u64 || nullstr (hour, nsid) /// - val: u64 || HLL (count (not cursor), estimator) /// +/// - Hourly record count ranking +/// - key: "hourly_rank_records" || u64 || u64 || nullstr (hour, count, nsid) +/// - val: [empty] +/// +/// - Hourly did estimate ranking +/// - key: "hourly_rank_dids" || u64 || u64 || nullstr (hour, dids estimate, nsid) +/// - val: [empty] +/// +/// /// - Weekly total record counts and dids estimate per collection -/// - key: "weekly_counts" || u64 || nullstr (hour, nsid) +/// - key: "weekly_counts" || u64 || nullstr (week, nsid) /// - val: u64 || HLL (count (not cursor), estimator) /// +/// - Weekly record count ranking +/// - key: "weekly_rank_records" || u64 || u64 || nullstr (week, count, nsid) +/// - val: [empty] +/// +/// - Weekly did estimate ranking +/// - key: "weekly_rank_dids" || u64 || u64 || nullstr (week, dids estimate, nsid) +/// - val: [empty] +/// +/// /// - All-time total record counts and dids estimate per collection /// - key: "ever_counts" || nullstr (nsid) /// - val: u64 || HLL (count (not cursor), estimator) /// -/// - TODO: sorted indexes for all-times? +/// - All-time total record record count ranking +/// - key: "ever_rank_records" || u64 || nullstr (count, nsid) +/// - val: [empty] +/// +/// - All-time did estimate ranking +/// - key: "ever_rank_dids" || u64 || nullstr (dids estimate, nsid) +/// - val: [empty] /// /// /// Partition: 'queues' diff --git a/ufos/src/store_types.rs b/ufos/src/store_types.rs index 2e9bdb8..2076de4 100644 --- a/ufos/src/store_types.rs +++ b/ufos/src/store_types.rs @@ -6,6 +6,18 @@ use bincode::{Decode, Encode}; use cardinality_estimator_safe::CardinalityEstimator; use std::ops::Range; +macro_rules! static_str { + ($prefix:expr, $name:ident) => { + #[derive(Debug, PartialEq)] + pub struct $name {} + impl StaticStr for $name { + fn static_str() -> &'static str { + $prefix + } + } + }; +} + /// key format: ["js_cursor"] #[derive(Debug, PartialEq)] pub struct JetstreamCursorKey {} @@ -289,6 +301,41 @@ impl DeleteAccountQueueKey { } pub type DeleteAccountQueueVal = Did; +/// big-endian encoded u64 for LSM prefix-fiendly key +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct KeyRank(u64); +impl DbBytes for KeyRank { + fn to_db_bytes(&self) -> Result, EncodingError> { + Ok(self.0.to_be_bytes().to_vec()) + } + fn from_db_bytes(bytes: &[u8]) -> Result<(Self, usize), EncodingError> { + if bytes.len() < 8 { + return Err(EncodingError::DecodeNotEnoughBytes); + } + let bytes8 = TryInto::<[u8; 8]>::try_into(&bytes[..8])?; + let rank = KeyRank(u64::from_be_bytes(bytes8)); + Ok((rank, 8)) + } +} + +pub type BucketedRankRecordsKey = + DbConcat, C>, KeyRank>, Nsid>; +impl BucketedRankRecordsKey +where + P: StaticStr + PartialEq + std::fmt::Debug, + C: DbBytes + PartialEq + std::fmt::Debug, +{ + pub fn new(cursor: C, rank: KeyRank, nsid: &Nsid) -> Self { + Self::from_pair( + DbConcat::from_pair(DbConcat::from_pair(Default::default(), cursor), rank), + nsid.clone(), + ) + } + pub fn update_rank(&mut self, new_rank: KeyRank) { + self.prefix.suffix = new_rank; + } +} + #[derive(Debug, PartialEq)] pub struct _HourlyRollupStaticStr {} impl StaticStr for _HourlyRollupStaticStr { @@ -308,6 +355,12 @@ impl HourlyRollupKey { } pub type HourlyRollupVal = CountsValue; +static_str!("hourly_rank_records", _HourlyRecordsStaticStr); +pub type HourlyRecordsKey = BucketedRankRecordsKey<_HourlyRecordsStaticStr, HourTruncatedCursor>; + +static_str!("hourly_rank_dids", _HourlyDidsStaticStr); +pub type HourlyDidsKey = BucketedRankRecordsKey<_HourlyDidsStaticStr, HourTruncatedCursor>; + #[derive(Debug, PartialEq)] pub struct _WeeklyRollupStaticStr {} impl StaticStr for _WeeklyRollupStaticStr { @@ -327,6 +380,12 @@ impl WeeklyRollupKey { } pub type WeeklyRollupVal = CountsValue; +static_str!("weekly_rank_records", _WeeklyRecordsStaticStr); +pub type WeeklyRecordsKey = BucketedRankRecordsKey<_WeeklyRecordsStaticStr, WeekTruncatedCursor>; + +static_str!("weekly_rank_dids", _WeeklyDidsStaticStr); +pub type WeeklyDidsKey = BucketedRankRecordsKey<_WeeklyDidsStaticStr, WeekTruncatedCursor>; + #[derive(Debug, PartialEq)] pub struct _AllTimeRollupStaticStr {} impl StaticStr for _AllTimeRollupStaticStr { @@ -346,6 +405,25 @@ impl AllTimeRollupKey { } pub type AllTimeRollupVal = CountsValue; +pub type AllTimeRankRecordsKey

= DbConcat, KeyRank>, Nsid>; +impl

AllTimeRankRecordsKey

+where + P: StaticStr + PartialEq + std::fmt::Debug, +{ + pub fn new(rank: KeyRank, nsid: &Nsid) -> Self { + Self::from_pair(DbConcat::from_pair(Default::default(), rank), nsid.clone()) + } + pub fn update_rank(&mut self, new_rank: KeyRank) { + self.prefix.suffix = new_rank; + } +} + +static_str!("ever_rank_records", _AllTimeRecordsStaticStr); +pub type AllTimeRecordsKey = AllTimeRankRecordsKey<_AllTimeRecordsStaticStr>; + +static_str!("ever_rank_dids", _AllTimeDidsStaticStr); +pub type AllTimeDidsKey = AllTimeRankRecordsKey<_AllTimeDidsStaticStr>; + #[derive(Debug, Copy, Clone, PartialEq, Hash, PartialOrd, Eq)] pub struct TruncatedCursor(u64); impl TruncatedCursor { -- 2.54.0 From f9075589ca20ddf5fa5d257bf08ce41bd34bc1f5 Mon Sep 17 00:00:00 2001 From: phil Date: Tue, 20 May 2025 12:31:07 -0400 Subject: [PATCH 2/7] use the static_str macro everywhere --- ufos/src/store_types.rs | 88 +++++++---------------------------------- 1 file changed, 14 insertions(+), 74 deletions(-) diff --git a/ufos/src/store_types.rs b/ufos/src/store_types.rs index 2076de4..c484dcb 100644 --- a/ufos/src/store_types.rs +++ b/ufos/src/store_types.rs @@ -18,35 +18,17 @@ macro_rules! static_str { }; } -/// key format: ["js_cursor"] -#[derive(Debug, PartialEq)] -pub struct JetstreamCursorKey {} -impl StaticStr for JetstreamCursorKey { - fn static_str() -> &'static str { - "js_cursor" - } -} +// key format: ["js_cursor"] +static_str!("js_cursor", JetstreamCursorKey); pub type JetstreamCursorValue = Cursor; -/// key format: ["rollup_cursor"] -#[derive(Debug, PartialEq)] -pub struct NewRollupCursorKey {} -impl StaticStr for NewRollupCursorKey { - fn static_str() -> &'static str { - "rollup_cursor" - } -} +// key format: ["rollup_cursor"] +static_str!("rollup_cursor", NewRollupCursorKey); // pub type NewRollupCursorKey = DbStaticStr<_NewRollupCursorKey>; /// value format: [rollup_cursor(Cursor)|collection(Nsid)] pub type NewRollupCursorValue = Cursor; -#[derive(Debug, PartialEq)] -pub struct _TrimCollectionStaticStr {} -impl StaticStr for _TrimCollectionStaticStr { - fn static_str() -> &'static str { - "trim_cursor" - } -} +static_str!("trim_cursor", _TrimCollectionStaticStr); type TrimCollectionCursorPrefix = DbStaticStr<_TrimCollectionStaticStr>; pub type TrimCollectionCursorKey = DbConcat; impl TrimCollectionCursorKey { @@ -56,24 +38,12 @@ impl TrimCollectionCursorKey { } pub type TrimCollectionCursorVal = Cursor; -/// key format: ["js_endpoint"] -#[derive(Debug, PartialEq)] -pub struct TakeoffKey {} -impl StaticStr for TakeoffKey { - fn static_str() -> &'static str { - "takeoff" - } -} +// key format: ["js_endpoint"] +static_str!("takeoff", TakeoffKey); pub type TakeoffValue = Cursor; -/// key format: ["js_endpoint"] -#[derive(Debug, PartialEq)] -pub struct JetstreamEndpointKey {} -impl StaticStr for JetstreamEndpointKey { - fn static_str() -> &'static str { - "js_endpoint" - } -} +// key format: ["js_endpoint"] +static_str!("js_endpoint", JetstreamEndpointKey); #[derive(Debug, PartialEq)] pub struct JetstreamEndpointValue(pub String); /// String wrapper for jetstream endpoint value @@ -199,13 +169,7 @@ impl From<(Cursor, &str, PutAction)> for RecordLocationVal { } } -#[derive(Debug, PartialEq)] -pub struct _LiveRecordsStaticStr {} -impl StaticStr for _LiveRecordsStaticStr { - fn static_str() -> &'static str { - "live_counts" - } -} +static_str!("live_counts", _LiveRecordsStaticStr); type LiveCountsStaticPrefix = DbStaticStr<_LiveRecordsStaticStr>; type LiveCountsCursorPrefix = DbConcat; @@ -285,13 +249,7 @@ impl Default for CountsValue { } } -#[derive(Debug, PartialEq)] -pub struct _DeleteAccountStaticStr {} -impl StaticStr for _DeleteAccountStaticStr { - fn static_str() -> &'static str { - "delete_acount" - } -} +static_str!("delete_acount", _DeleteAccountStaticStr); pub type DeleteAccountStaticPrefix = DbStaticStr<_DeleteAccountStaticStr>; pub type DeleteAccountQueueKey = DbConcat; impl DeleteAccountQueueKey { @@ -336,13 +294,7 @@ where } } -#[derive(Debug, PartialEq)] -pub struct _HourlyRollupStaticStr {} -impl StaticStr for _HourlyRollupStaticStr { - fn static_str() -> &'static str { - "hourly_counts" - } -} +static_str!("hourly_counts", _HourlyRollupStaticStr); pub type HourlyRollupStaticPrefix = DbStaticStr<_HourlyRollupStaticStr>; pub type HourlyRollupKey = DbConcat, Nsid>; impl HourlyRollupKey { @@ -361,13 +313,7 @@ pub type HourlyRecordsKey = BucketedRankRecordsKey<_HourlyRecordsStaticStr, Hour static_str!("hourly_rank_dids", _HourlyDidsStaticStr); pub type HourlyDidsKey = BucketedRankRecordsKey<_HourlyDidsStaticStr, HourTruncatedCursor>; -#[derive(Debug, PartialEq)] -pub struct _WeeklyRollupStaticStr {} -impl StaticStr for _WeeklyRollupStaticStr { - fn static_str() -> &'static str { - "weekly_counts" - } -} +static_str!("weekly_counts", _WeeklyRollupStaticStr); pub type WeeklyRollupStaticPrefix = DbStaticStr<_WeeklyRollupStaticStr>; pub type WeeklyRollupKey = DbConcat, Nsid>; impl WeeklyRollupKey { @@ -386,13 +332,7 @@ pub type WeeklyRecordsKey = BucketedRankRecordsKey<_WeeklyRecordsStaticStr, Week static_str!("weekly_rank_dids", _WeeklyDidsStaticStr); pub type WeeklyDidsKey = BucketedRankRecordsKey<_WeeklyDidsStaticStr, WeekTruncatedCursor>; -#[derive(Debug, PartialEq)] -pub struct _AllTimeRollupStaticStr {} -impl StaticStr for _AllTimeRollupStaticStr { - fn static_str() -> &'static str { - "ever_counts" - } -} +static_str!("ever_counts", _AllTimeRollupStaticStr); pub type AllTimeRollupStaticPrefix = DbStaticStr<_AllTimeRollupStaticStr>; pub type AllTimeRollupKey = DbConcat; impl AllTimeRollupKey { -- 2.54.0 From 4e1ea47d18701fd18278131ef6a0e37720e1f2b4 Mon Sep 17 00:00:00 2001 From: phil Date: Tue, 20 May 2025 12:42:02 -0400 Subject: [PATCH 3/7] remove old debug code --- ufos/src/storage_fjall.rs | 21 +++------------------ 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/ufos/src/storage_fjall.rs b/ufos/src/storage_fjall.rs index 3e0c464..584b56b 100644 --- a/ufos/src/storage_fjall.rs +++ b/ufos/src/storage_fjall.rs @@ -600,17 +600,12 @@ impl FjallWriter { for ((nsid, rollup), counts) in counts_by_rollup { let key_bytes = match rollup { Rollup::Hourly(hourly_cursor) => { - let k = HourlyRollupKey::new(hourly_cursor, &nsid); - k.to_db_bytes()? + HourlyRollupKey::new(hourly_cursor, &nsid).to_db_bytes()? } Rollup::Weekly(weekly_cursor) => { - let k = WeeklyRollupKey::new(weekly_cursor, &nsid); - k.to_db_bytes()? - } - Rollup::AllTime => { - let k = AllTimeRollupKey::new(&nsid); - k.to_db_bytes()? + WeeklyRollupKey::new(weekly_cursor, &nsid).to_db_bytes()? } + Rollup::AllTime => AllTimeRollupKey::new(&nsid).to_db_bytes()?, }; let mut rolled: CountsValue = self .rollups @@ -620,16 +615,6 @@ impl FjallWriter { .transpose()? .unwrap_or_default(); - // try to round-trip before inserting, for funsies - let tripppin = counts.to_db_bytes()?; - let (and_back, n) = CountsValue::from_db_bytes(&tripppin)?; - assert_eq!(n, tripppin.len()); - assert_eq!(counts.prefix, and_back.prefix); - assert_eq!(counts.dids().estimate(), and_back.dids().estimate()); - if counts.records() > 200_000_000_000 { - panic!("COUNTS maybe wtf? {counts:?}") - } - rolled.merge(&counts); batch.insert(&self.rollups, &key_bytes, &rolled.to_db_bytes()?); } -- 2.54.0 From a04e8a50e70205a5834901ead5ccade055cb09c3 Mon Sep 17 00:00:00 2001 From: phil Date: Tue, 20 May 2025 13:03:08 -0400 Subject: [PATCH 4/7] save ranks --- ufos/src/storage_fjall.rs | 90 +++++++++++++++++++++++++++++++++++---- ufos/src/store_types.rs | 10 +++++ 2 files changed, 91 insertions(+), 9 deletions(-) diff --git a/ufos/src/storage_fjall.rs b/ufos/src/storage_fjall.rs index 584b56b..d9961d5 100644 --- a/ufos/src/storage_fjall.rs +++ b/ufos/src/storage_fjall.rs @@ -2,12 +2,13 @@ use crate::db_types::{db_complete, DbBytes, DbStaticStr, StaticStr}; use crate::error::StorageError; use crate::storage::{StorageResult, StorageWhatever, StoreBackground, StoreReader, StoreWriter}; use crate::store_types::{ - AllTimeRollupKey, CountsValue, DeleteAccountQueueKey, DeleteAccountQueueVal, - HourTruncatedCursor, HourlyRollupKey, JetstreamCursorKey, JetstreamCursorValue, - JetstreamEndpointKey, JetstreamEndpointValue, LiveCountsKey, NewRollupCursorKey, - NewRollupCursorValue, NsidRecordFeedKey, NsidRecordFeedVal, RecordLocationKey, - RecordLocationMeta, RecordLocationVal, RecordRawValue, TakeoffKey, TakeoffValue, - TrimCollectionCursorKey, WeekTruncatedCursor, WeeklyRollupKey, + AllTimeDidsKey, AllTimeRecordsKey, AllTimeRollupKey, CountsValue, DeleteAccountQueueKey, + DeleteAccountQueueVal, HourTruncatedCursor, HourlyDidsKey, HourlyRecordsKey, HourlyRollupKey, + JetstreamCursorKey, JetstreamCursorValue, JetstreamEndpointKey, JetstreamEndpointValue, + LiveCountsKey, NewRollupCursorKey, NewRollupCursorValue, NsidRecordFeedKey, NsidRecordFeedVal, + RecordLocationKey, RecordLocationMeta, RecordLocationVal, RecordRawValue, TakeoffKey, + TakeoffValue, TrimCollectionCursorKey, WeekTruncatedCursor, WeeklyDidsKey, WeeklyRecordsKey, + WeeklyRollupKey, }; use crate::{CommitAction, ConsumerInfo, Did, EventBatch, Nsid, TopCollections, UFOsRecord}; use async_trait::async_trait; @@ -597,8 +598,9 @@ impl FjallWriter { last_cursor = key.cursor(); } + // go through each new rollup thing and merge it with whatever might already be in the db for ((nsid, rollup), counts) in counts_by_rollup { - let key_bytes = match rollup { + let rollup_key_bytes = match rollup { Rollup::Hourly(hourly_cursor) => { HourlyRollupKey::new(hourly_cursor, &nsid).to_db_bytes()? } @@ -609,14 +611,84 @@ impl FjallWriter { }; let mut rolled: CountsValue = self .rollups - .get(&key_bytes)? + .get(&rollup_key_bytes)? .as_deref() .map(db_complete::) .transpose()? .unwrap_or_default(); + // now that we have values, we can fetch the ranks + // TODO: save .records() and .estimate() and then do this match only once later + let (rank_records_key_bytes, rank_dids_key_bytes) = match rollup { + Rollup::Hourly(hourly_cursor) => ( + HourlyRecordsKey::new(hourly_cursor, rolled.records().into(), &nsid) + .to_db_bytes()?, + HourlyDidsKey::new( + hourly_cursor, + (rolled.dids().estimate() as u64).into(), + &nsid, + ) + .to_db_bytes()?, + ), + Rollup::Weekly(weekly_cursor) => ( + WeeklyRecordsKey::new(weekly_cursor, rolled.records().into(), &nsid) + .to_db_bytes()?, + WeeklyDidsKey::new( + weekly_cursor, + (rolled.dids().estimate() as u64).into(), + &nsid, + ) + .to_db_bytes()?, + ), + Rollup::AllTime => ( + AllTimeRecordsKey::new(rolled.records().into(), &nsid).to_db_bytes()?, + AllTimeDidsKey::new((rolled.dids().estimate() as u64).into(), &nsid) + .to_db_bytes()?, + ), + }; + + // update the rollup rolled.merge(&counts); - batch.insert(&self.rollups, &key_bytes, &rolled.to_db_bytes()?); + + // replace rank entries + let (new_rank_records_key_bytes, new_rank_dids_key_bytes) = match rollup { + Rollup::Hourly(hourly_cursor) => ( + HourlyRecordsKey::new(hourly_cursor, rolled.records().into(), &nsid) + .to_db_bytes()?, + HourlyDidsKey::new( + hourly_cursor, + (rolled.dids().estimate() as u64).into(), + &nsid, + ) + .to_db_bytes()?, + ), + Rollup::Weekly(weekly_cursor) => ( + WeeklyRecordsKey::new(weekly_cursor, rolled.records().into(), &nsid) + .to_db_bytes()?, + WeeklyDidsKey::new( + weekly_cursor, + (rolled.dids().estimate() as u64).into(), + &nsid, + ) + .to_db_bytes()?, + ), + Rollup::AllTime => ( + AllTimeRecordsKey::new(rolled.records().into(), &nsid).to_db_bytes()?, + AllTimeDidsKey::new((rolled.dids().estimate() as u64).into(), &nsid) + .to_db_bytes()?, + ), + }; + if new_rank_records_key_bytes != rank_records_key_bytes { + batch.remove(&self.rollups, &rank_records_key_bytes); + batch.insert(&self.rollups, &new_rank_records_key_bytes, ""); + } + if new_rank_dids_key_bytes != rank_dids_key_bytes { + batch.remove(&self.rollups, &rank_dids_key_bytes); + batch.insert(&self.rollups, &new_rank_dids_key_bytes, ""); + } + + // set the updated rollup + batch.insert(&self.rollups, &rollup_key_bytes, &rolled.to_db_bytes()?); } insert_batch_static_neu::(&mut batch, &self.global, last_cursor)?; diff --git a/ufos/src/store_types.rs b/ufos/src/store_types.rs index c484dcb..780ed46 100644 --- a/ufos/src/store_types.rs +++ b/ufos/src/store_types.rs @@ -275,6 +275,16 @@ impl DbBytes for KeyRank { Ok((rank, 8)) } } +impl From for KeyRank { + fn from(n: u64) -> Self { + Self(n) + } +} +impl From for u64 { + fn from(kr: KeyRank) -> Self { + kr.0 + } +} pub type BucketedRankRecordsKey = DbConcat, C>, KeyRank>, Nsid>; -- 2.54.0 From 66b77a23b4cd9a7cae5d734e69e9d17ce3b9915d Mon Sep 17 00:00:00 2001 From: phil Date: Tue, 20 May 2025 13:23:48 -0400 Subject: [PATCH 5/7] simplify(??) the rank update stuff --- ufos/src/storage_fjall.rs | 118 ++++++++++++++++++-------------------- ufos/src/store_types.rs | 10 ++-- 2 files changed, 60 insertions(+), 68 deletions(-) diff --git a/ufos/src/storage_fjall.rs b/ufos/src/storage_fjall.rs index d9961d5..ef3fa51 100644 --- a/ufos/src/storage_fjall.rs +++ b/ufos/src/storage_fjall.rs @@ -617,77 +617,69 @@ impl FjallWriter { .transpose()? .unwrap_or_default(); - // now that we have values, we can fetch the ranks - // TODO: save .records() and .estimate() and then do this match only once later - let (rank_records_key_bytes, rank_dids_key_bytes) = match rollup { - Rollup::Hourly(hourly_cursor) => ( - HourlyRecordsKey::new(hourly_cursor, rolled.records().into(), &nsid) - .to_db_bytes()?, - HourlyDidsKey::new( - hourly_cursor, - (rolled.dids().estimate() as u64).into(), - &nsid, - ) - .to_db_bytes()?, - ), - Rollup::Weekly(weekly_cursor) => ( - WeeklyRecordsKey::new(weekly_cursor, rolled.records().into(), &nsid) - .to_db_bytes()?, - WeeklyDidsKey::new( - weekly_cursor, - (rolled.dids().estimate() as u64).into(), - &nsid, - ) - .to_db_bytes()?, - ), - Rollup::AllTime => ( - AllTimeRecordsKey::new(rolled.records().into(), &nsid).to_db_bytes()?, - AllTimeDidsKey::new((rolled.dids().estimate() as u64).into(), &nsid) - .to_db_bytes()?, - ), - }; + // now that we have values, we can know the exising ranks + let before_records_count = rolled.records(); + let before_dids_estimate = rolled.dids().estimate() as u64; // update the rollup rolled.merge(&counts); // replace rank entries - let (new_rank_records_key_bytes, new_rank_dids_key_bytes) = match rollup { - Rollup::Hourly(hourly_cursor) => ( - HourlyRecordsKey::new(hourly_cursor, rolled.records().into(), &nsid) - .to_db_bytes()?, - HourlyDidsKey::new( - hourly_cursor, - (rolled.dids().estimate() as u64).into(), - &nsid, - ) - .to_db_bytes()?, - ), - Rollup::Weekly(weekly_cursor) => ( - WeeklyRecordsKey::new(weekly_cursor, rolled.records().into(), &nsid) - .to_db_bytes()?, - WeeklyDidsKey::new( - weekly_cursor, - (rolled.dids().estimate() as u64).into(), - &nsid, - ) - .to_db_bytes()?, - ), - Rollup::AllTime => ( - AllTimeRecordsKey::new(rolled.records().into(), &nsid).to_db_bytes()?, - AllTimeDidsKey::new((rolled.dids().estimate() as u64).into(), &nsid) - .to_db_bytes()?, - ), + let (old_records, new_records, dids) = match rollup { + Rollup::Hourly(hourly_cursor) => { + let old_records = + HourlyRecordsKey::new(hourly_cursor, before_records_count.into(), &nsid); + let new_records = old_records.with_rank(rolled.records().into()); + let new_estimate = rolled.dids().estimate() as u64; + let dids = if new_estimate == before_dids_estimate { + None + } else { + let old_dids = + HourlyDidsKey::new(hourly_cursor, before_dids_estimate.into(), &nsid); + let new_dids = old_dids.with_rank(new_estimate.into()); + Some((old_dids.to_db_bytes()?, new_dids.to_db_bytes()?)) + }; + (old_records.to_db_bytes()?, new_records.to_db_bytes()?, dids) + } + Rollup::Weekly(weekly_cursor) => { + let old_records = + WeeklyRecordsKey::new(weekly_cursor, before_records_count.into(), &nsid); + let new_records = old_records.with_rank(rolled.records().into()); + let new_estimate = rolled.dids().estimate() as u64; + let dids = if new_estimate == before_dids_estimate { + None + } else { + let old_dids = + WeeklyDidsKey::new(weekly_cursor, before_dids_estimate.into(), &nsid); + let new_dids = old_dids.with_rank(new_estimate.into()); + Some((old_dids.to_db_bytes()?, new_dids.to_db_bytes()?)) + }; + (old_records.to_db_bytes()?, new_records.to_db_bytes()?, dids) + } + Rollup::AllTime => { + let old_records = AllTimeRecordsKey::new(before_records_count.into(), &nsid); + let new_records = old_records.with_rank(rolled.records().into()); + let new_estimate = rolled.dids().estimate() as u64; + let dids = if new_estimate == before_dids_estimate { + None + } else { + let old_dids = AllTimeDidsKey::new(before_dids_estimate.into(), &nsid); + let new_dids = old_dids.with_rank(new_estimate.into()); + Some((old_dids.to_db_bytes()?, new_dids.to_db_bytes()?)) + }; + (old_records.to_db_bytes()?, new_records.to_db_bytes()?, dids) + } }; - if new_rank_records_key_bytes != rank_records_key_bytes { - batch.remove(&self.rollups, &rank_records_key_bytes); - batch.insert(&self.rollups, &new_rank_records_key_bytes, ""); - } - if new_rank_dids_key_bytes != rank_dids_key_bytes { - batch.remove(&self.rollups, &rank_dids_key_bytes); - batch.insert(&self.rollups, &new_rank_dids_key_bytes, ""); + + // replace the ranks + batch.remove(&self.rollups, &old_records); + batch.insert(&self.rollups, &new_records, ""); + if let Some((old_dids, new_dids)) = dids { + batch.remove(&self.rollups, &old_dids); + batch.insert(&self.rollups, &new_dids, ""); } - // set the updated rollup + // replace the rollup batch.insert(&self.rollups, &rollup_key_bytes, &rolled.to_db_bytes()?); } diff --git a/ufos/src/store_types.rs b/ufos/src/store_types.rs index 780ed46..6ad151e 100644 --- a/ufos/src/store_types.rs +++ b/ufos/src/store_types.rs @@ -291,7 +291,7 @@ pub type BucketedRankRecordsKey = impl BucketedRankRecordsKey where P: StaticStr + PartialEq + std::fmt::Debug, - C: DbBytes + PartialEq + std::fmt::Debug, + C: DbBytes + PartialEq + std::fmt::Debug + Clone, { pub fn new(cursor: C, rank: KeyRank, nsid: &Nsid) -> Self { Self::from_pair( @@ -299,8 +299,8 @@ where nsid.clone(), ) } - pub fn update_rank(&mut self, new_rank: KeyRank) { - self.prefix.suffix = new_rank; + pub fn with_rank(&self, new_rank: KeyRank) -> Self { + Self::new(self.prefix.prefix.suffix.clone(), new_rank, &self.suffix) } } @@ -363,8 +363,8 @@ where pub fn new(rank: KeyRank, nsid: &Nsid) -> Self { Self::from_pair(DbConcat::from_pair(Default::default(), rank), nsid.clone()) } - pub fn update_rank(&mut self, new_rank: KeyRank) { - self.prefix.suffix = new_rank; + pub fn with_rank(&self, new_rank: KeyRank) -> Self { + Self::new(new_rank, &self.suffix) } } -- 2.54.0 From c47b582e470498f4766f63e4142ad5e1d62a53e0 Mon Sep 17 00:00:00 2001 From: phil Date: Tue, 20 May 2025 23:07:47 -0400 Subject: [PATCH 6/7] get top collections by records count --- ufos/src/lib.rs | 24 ++++++++++++++++++++ ufos/src/server.rs | 18 ++++++++++++++- ufos/src/storage.rs | 11 +++++++++- ufos/src/storage_fjall.rs | 46 ++++++++++++++++++++++++++++++++++++++- ufos/src/storage_mem.rs | 21 ++++++++++++------ ufos/src/store_types.rs | 20 +++++++++++------ 6 files changed, 123 insertions(+), 17 deletions(-) diff --git a/ufos/src/lib.rs b/ufos/src/lib.rs index eec9685..52d893c 100644 --- a/ufos/src/lib.rs +++ b/ufos/src/lib.rs @@ -249,6 +249,30 @@ impl From for Vec { } } +#[derive(Debug)] +pub struct QueryPeriod { + from: Option, + until: Option, +} +impl QueryPeriod { + pub fn all_time() -> Self { + QueryPeriod { + from: None, + until: None, + } + } + pub fn is_all_time(&self) -> bool { + self.from.is_none() && self.until.is_none() + } +} + +#[derive(Debug, Serialize, JsonSchema)] +pub struct Count { + thing: String, + records: u64, + dids_estimate: u64, +} + #[cfg(test)] mod tests { use super::*; diff --git a/ufos/src/server.rs b/ufos/src/server.rs index c86d679..5ae6163 100644 --- a/ufos/src/server.rs +++ b/ufos/src/server.rs @@ -1,6 +1,6 @@ use crate::index_html::INDEX_HTML; use crate::storage::StoreReader; -use crate::{ConsumerInfo, Nsid, TopCollections, UFOsRecord}; +use crate::{ConsumerInfo, Count, Nsid, QueryPeriod, TopCollections, UFOsRecord}; use dropshot::endpoint; use dropshot::ApiDescription; use dropshot::Body; @@ -213,6 +213,21 @@ async fn get_records_total_seen( ok_cors(seen_by_collection) } +/// Get top collections by count +#[endpoint { + method = GET, + path = "/collections/by-count" +}] +async fn get_top_collections_by_count(ctx: RequestContext) -> OkCorsResponse> { + let Context { storage, .. } = ctx.context(); + let collections = storage + .get_top_collections_by_count(100, QueryPeriod::all_time()) + .await + .map_err(|e| HttpError::for_internal_error(format!("oh shoot: {e:?}")))?; + + ok_cors(collections) +} + /// Get top collections /// /// The format of this API response will be changing soon. @@ -244,6 +259,7 @@ pub async fn serve(storage: impl StoreReader + 'static) -> Result<(), String> { api.register(get_meta_info).unwrap(); api.register(get_records_by_collections).unwrap(); api.register(get_records_total_seen).unwrap(); + api.register(get_top_collections_by_count).unwrap(); api.register(get_top_collections).unwrap(); let context = Context { diff --git a/ufos/src/storage.rs b/ufos/src/storage.rs index ffd0599..5839da7 100644 --- a/ufos/src/storage.rs +++ b/ufos/src/storage.rs @@ -1,4 +1,7 @@ -use crate::{error::StorageError, ConsumerInfo, Cursor, EventBatch, TopCollections, UFOsRecord}; +use crate::{ + error::StorageError, ConsumerInfo, Count, Cursor, EventBatch, QueryPeriod, TopCollections, + UFOsRecord, +}; use async_trait::async_trait; use jetstream::exports::{Did, Nsid}; use std::collections::HashSet; @@ -67,6 +70,12 @@ pub trait StoreReader: Send + Sync { async fn get_consumer_info(&self) -> StorageResult; + async fn get_top_collections_by_count( + &self, + limit: usize, + period: QueryPeriod, + ) -> StorageResult>; + async fn get_top_collections(&self) -> StorageResult; async fn get_counts_by_collection(&self, collection: &Nsid) -> StorageResult<(u64, u64)>; diff --git a/ufos/src/storage_fjall.rs b/ufos/src/storage_fjall.rs index ef3fa51..01a956c 100644 --- a/ufos/src/storage_fjall.rs +++ b/ufos/src/storage_fjall.rs @@ -10,7 +10,10 @@ use crate::store_types::{ TakeoffValue, TrimCollectionCursorKey, WeekTruncatedCursor, WeeklyDidsKey, WeeklyRecordsKey, WeeklyRollupKey, }; -use crate::{CommitAction, ConsumerInfo, Did, EventBatch, Nsid, TopCollections, UFOsRecord}; +use crate::{ + CommitAction, ConsumerInfo, Count, Did, EventBatch, Nsid, QueryPeriod, TopCollections, + UFOsRecord, +}; use async_trait::async_trait; use fjall::{Batch as FjallBatch, Config, Keyspace, PartitionCreateOptions, PartitionHandle}; use jetstream::events::Cursor; @@ -339,6 +342,36 @@ impl FjallReader { }) } + fn get_top_collections_by_count( + &self, + limit: usize, + period: QueryPeriod, + ) -> StorageResult> { + Ok(if period.is_all_time() { + let snapshot = self.rollups.snapshot(); + let mut out = Vec::with_capacity(limit); + let prefix = AllTimeRecordsKey::from_prefix_to_db_bytes(&Default::default())?; + for kv in snapshot.prefix(prefix).rev().take(limit) { + let (key_bytes, _) = kv?; + let key = db_complete::(&key_bytes)?; + let rollup_key = AllTimeRollupKey::new(key.collection()); + let db_count_bytes = snapshot.get(rollup_key.to_db_bytes()?)?.expect( + "integrity: all-time rank rollup must have corresponding all-time count rollup", + ); + let db_counts = db_complete::(&db_count_bytes)?; + assert_eq!(db_counts.records(), key.records()); + out.push(Count { + thing: key.collection().to_string(), + records: db_counts.records(), + dids_estimate: db_counts.dids().estimate() as u64, + }); + } + out + } else { + todo!() + }) + } + fn get_top_collections(&self) -> Result { // TODO: limit nsid traversal depth // TODO: limit nsid traversal breadth @@ -480,6 +513,17 @@ impl StoreReader for FjallReader { let s = self.clone(); tokio::task::spawn_blocking(move || FjallReader::get_consumer_info(&s)).await? } + async fn get_top_collections_by_count( + &self, + limit: usize, + period: QueryPeriod, + ) -> StorageResult> { + let s = self.clone(); + tokio::task::spawn_blocking(move || { + FjallReader::get_top_collections_by_count(&s, limit, period) + }) + .await? + } async fn get_top_collections(&self) -> Result { let s = self.clone(); tokio::task::spawn_blocking(move || FjallReader::get_top_collections(&s)).await? diff --git a/ufos/src/storage_mem.rs b/ufos/src/storage_mem.rs index 2afece8..d798b50 100644 --- a/ufos/src/storage_mem.rs +++ b/ufos/src/storage_mem.rs @@ -12,16 +12,16 @@ use crate::store_types::{ RecordLocationMeta, RecordLocationVal, RecordRawValue, TakeoffKey, TakeoffValue, WeekTruncatedCursor, WeeklyRollupKey, }; -use crate::{CommitAction, ConsumerInfo, Did, EventBatch, Nsid, TopCollections, UFOsRecord}; +use crate::{ + CommitAction, ConsumerInfo, Count, Did, EventBatch, Nsid, QueryPeriod, TopCollections, + UFOsRecord, +}; use async_trait::async_trait; use jetstream::events::Cursor; use lsm_tree::range::prefix_to_range; -use std::collections::BTreeMap; -use std::collections::HashMap; -use std::collections::HashSet; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::path::Path; -use std::sync::Mutex; -use std::sync::RwLock; +use std::sync::{Mutex, RwLock}; use std::time::SystemTime; const MAX_BATCHED_CLEANUP_SIZE: usize = 1024; // try to commit progress for longer feeds @@ -584,10 +584,17 @@ impl StoreReader for MemReader { let s = self.clone(); tokio::task::spawn_blocking(move || MemReader::get_consumer_info(&s)).await? } - async fn get_top_collections(&self) -> Result { + async fn get_top_collections(&self) -> StorageResult { let s = self.clone(); tokio::task::spawn_blocking(move || MemReader::get_top_collections(&s)).await? } + async fn get_top_collections_by_count( + &self, + _: usize, + _: QueryPeriod, + ) -> StorageResult> { + todo!() + } async fn get_counts_by_collection(&self, collection: &Nsid) -> StorageResult<(u64, u64)> { let s = self.clone(); let collection = collection.clone(); diff --git a/ufos/src/store_types.rs b/ufos/src/store_types.rs index 6ad151e..8b5c183 100644 --- a/ufos/src/store_types.rs +++ b/ufos/src/store_types.rs @@ -287,7 +287,7 @@ impl From for u64 { } pub type BucketedRankRecordsKey = - DbConcat, C>, KeyRank>, Nsid>; + DbConcat, C>, DbConcat>; impl BucketedRankRecordsKey where P: StaticStr + PartialEq + std::fmt::Debug, @@ -295,12 +295,12 @@ where { pub fn new(cursor: C, rank: KeyRank, nsid: &Nsid) -> Self { Self::from_pair( - DbConcat::from_pair(DbConcat::from_pair(Default::default(), cursor), rank), - nsid.clone(), + DbConcat::from_pair(Default::default(), cursor), + DbConcat::from_pair(rank, nsid.clone()), ) } pub fn with_rank(&self, new_rank: KeyRank) -> Self { - Self::new(self.prefix.prefix.suffix.clone(), new_rank, &self.suffix) + Self::new(self.prefix.suffix.clone(), new_rank, &self.suffix.suffix) } } @@ -355,16 +355,22 @@ impl AllTimeRollupKey { } pub type AllTimeRollupVal = CountsValue; -pub type AllTimeRankRecordsKey

= DbConcat, KeyRank>, Nsid>; +pub type AllTimeRankRecordsKey

= DbConcat, DbConcat>; impl

AllTimeRankRecordsKey

where P: StaticStr + PartialEq + std::fmt::Debug, { pub fn new(rank: KeyRank, nsid: &Nsid) -> Self { - Self::from_pair(DbConcat::from_pair(Default::default(), rank), nsid.clone()) + Self::from_pair(Default::default(), DbConcat::from_pair(rank, nsid.clone())) } pub fn with_rank(&self, new_rank: KeyRank) -> Self { - Self::new(new_rank, &self.suffix) + Self::new(new_rank, &self.suffix.suffix) + } + pub fn records(&self) -> u64 { + self.suffix.prefix.0 + } + pub fn collection(&self) -> &Nsid { + &self.suffix.suffix } } -- 2.54.0 From f9c70c293a2c3cbbd467b1b4190ea1233c769518 Mon Sep 17 00:00:00 2001 From: phil Date: Wed, 21 May 2025 10:47:38 -0400 Subject: [PATCH 7/7] get top collections by uniq dids estimate --- ufos/src/server.rs | 18 +++++++++++++++- ufos/src/storage.rs | 6 ++++++ ufos/src/storage_fjall.rs | 43 ++++++++++++++++++++++++++++++++++++++- ufos/src/storage_mem.rs | 7 +++++++ ufos/src/store_types.rs | 2 +- 5 files changed, 73 insertions(+), 3 deletions(-) diff --git a/ufos/src/server.rs b/ufos/src/server.rs index 5ae6163..3ad04c6 100644 --- a/ufos/src/server.rs +++ b/ufos/src/server.rs @@ -213,7 +213,7 @@ async fn get_records_total_seen( ok_cors(seen_by_collection) } -/// Get top collections by count +/// Get top collections by record count #[endpoint { method = GET, path = "/collections/by-count" @@ -228,6 +228,21 @@ async fn get_top_collections_by_count(ctx: RequestContext) -> OkCorsRes ok_cors(collections) } +/// Get top collections by estimated unique DIDs +#[endpoint { + method = GET, + path = "/collections/by-dids" +}] +async fn get_top_collections_by_dids(ctx: RequestContext) -> OkCorsResponse> { + let Context { storage, .. } = ctx.context(); + let collections = storage + .get_top_collections_by_dids(100, QueryPeriod::all_time()) + .await + .map_err(|e| HttpError::for_internal_error(format!("oh shoot: {e:?}")))?; + + ok_cors(collections) +} + /// Get top collections /// /// The format of this API response will be changing soon. @@ -260,6 +275,7 @@ pub async fn serve(storage: impl StoreReader + 'static) -> Result<(), String> { api.register(get_records_by_collections).unwrap(); api.register(get_records_total_seen).unwrap(); api.register(get_top_collections_by_count).unwrap(); + api.register(get_top_collections_by_dids).unwrap(); api.register(get_top_collections).unwrap(); let context = Context { diff --git a/ufos/src/storage.rs b/ufos/src/storage.rs index 5839da7..103d1b5 100644 --- a/ufos/src/storage.rs +++ b/ufos/src/storage.rs @@ -76,6 +76,12 @@ pub trait StoreReader: Send + Sync { period: QueryPeriod, ) -> StorageResult>; + async fn get_top_collections_by_dids( + &self, + limit: usize, + period: QueryPeriod, + ) -> StorageResult>; + async fn get_top_collections(&self) -> StorageResult; async fn get_counts_by_collection(&self, collection: &Nsid) -> StorageResult<(u64, u64)>; diff --git a/ufos/src/storage_fjall.rs b/ufos/src/storage_fjall.rs index 01a956c..dc2622d 100644 --- a/ufos/src/storage_fjall.rs +++ b/ufos/src/storage_fjall.rs @@ -359,7 +359,37 @@ impl FjallReader { "integrity: all-time rank rollup must have corresponding all-time count rollup", ); let db_counts = db_complete::(&db_count_bytes)?; - assert_eq!(db_counts.records(), key.records()); + assert_eq!(db_counts.records(), key.count()); + out.push(Count { + thing: key.collection().to_string(), + records: db_counts.records(), + dids_estimate: db_counts.dids().estimate() as u64, + }); + } + out + } else { + todo!() + }) + } + + fn get_top_collections_by_dids( + &self, + limit: usize, + period: QueryPeriod, + ) -> StorageResult> { + Ok(if period.is_all_time() { + let snapshot = self.rollups.snapshot(); + let mut out = Vec::with_capacity(limit); + let prefix = AllTimeDidsKey::from_prefix_to_db_bytes(&Default::default())?; + for kv in snapshot.prefix(prefix).rev().take(limit) { + let (key_bytes, _) = kv?; + let key = db_complete::(&key_bytes)?; + let rollup_key = AllTimeRollupKey::new(key.collection()); + let db_count_bytes = snapshot.get(rollup_key.to_db_bytes()?)?.expect( + "integrity: all-time rank rollup must have corresponding all-time count rollup", + ); + let db_counts = db_complete::(&db_count_bytes)?; + assert_eq!(db_counts.dids().estimate() as u64, key.count()); out.push(Count { thing: key.collection().to_string(), records: db_counts.records(), @@ -524,6 +554,17 @@ impl StoreReader for FjallReader { }) .await? } + async fn get_top_collections_by_dids( + &self, + limit: usize, + period: QueryPeriod, + ) -> StorageResult> { + let s = self.clone(); + tokio::task::spawn_blocking(move || { + FjallReader::get_top_collections_by_dids(&s, limit, period) + }) + .await? + } async fn get_top_collections(&self) -> Result { let s = self.clone(); tokio::task::spawn_blocking(move || FjallReader::get_top_collections(&s)).await? diff --git a/ufos/src/storage_mem.rs b/ufos/src/storage_mem.rs index d798b50..23e41ee 100644 --- a/ufos/src/storage_mem.rs +++ b/ufos/src/storage_mem.rs @@ -595,6 +595,13 @@ impl StoreReader for MemReader { ) -> StorageResult> { todo!() } + async fn get_top_collections_by_dids( + &self, + _: usize, + _: QueryPeriod, + ) -> StorageResult> { + todo!() + } async fn get_counts_by_collection(&self, collection: &Nsid) -> StorageResult<(u64, u64)> { let s = self.clone(); let collection = collection.clone(); diff --git a/ufos/src/store_types.rs b/ufos/src/store_types.rs index 8b5c183..359856f 100644 --- a/ufos/src/store_types.rs +++ b/ufos/src/store_types.rs @@ -366,7 +366,7 @@ where pub fn with_rank(&self, new_rank: KeyRank) -> Self { Self::new(new_rank, &self.suffix.suffix) } - pub fn records(&self) -> u64 { + pub fn count(&self) -> u64 { self.suffix.prefix.0 } pub fn collection(&self) -> &Nsid {