From c1614c91e2582b8607eda16a9935a29030960485 Mon Sep 17 00:00:00 2001 From: phil Date: Tue, 20 May 2025 12:25:28 -0400 Subject: [PATCH 1/7] half a dozen new indices for ranking could probably get away with either just ranking records (upper bound for DIDs) or DIDs (lower bound for records) but ordering by the other would probably be messy / approximate. these indices should be pretty cheap to store, so doing it is must a matter of overcoming my laziness to that effect i'm finally adding a macro for the static_str thing, and made the keys generic-ish. --- ufos/src/storage_fjall.rs | 29 ++++++++++++++- ufos/src/store_types.rs | 78 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 105 insertions(+), 2 deletions(-) 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 { -- 2.54.0