Merge pull request #22 from at-microcosm/collection-stats

collection stats

Author: uniphil

Cargo.lock

@@ -23,6 +23,7 @@ schemars = { version = "0.8.22", features = ["raw_value", "chrono"] }
 semver = "1.0.26"
 serde = "1.0.219"
 serde_json = "1.0.140"
+serde_qs = "1.0.0-rc.3"
 sha2 = "0.10.9"
 thiserror = "2.0.12"
 tokio = { version = "1.44.2", features = ["full", "sync", "time"] }

ufos/Cargo.toml

@@ -6,7 +6,6 @@ pub mod index_html;
 pub mod server;
 pub mod storage;
 pub mod storage_fjall;
-pub mod storage_mem;
 pub mod store_types;
 
 use crate::error::BatchInsertError;
@@ -287,6 +286,8 @@ pub struct NsidCount {
 #[derive(Debug, Serialize, JsonSchema)]
 pub struct JustCount {
     creates: u64,
+    updates: u64,
+    deletes: u64,
     dids_estimate: u64,
 }
 

ufos/src/lib.rs

@@ -7,7 +7,6 @@ use ufos::file_consumer;
 use ufos::server;
 use ufos::storage::{StorageWhatever, StoreBackground, StoreReader, StoreWriter};
 use ufos::storage_fjall::FjallStorage;
-use ufos::storage_mem::MemStorage;
 use ufos::store_types::SketchSecretPrefix;
 use ufos::{nice_duration, ConsumerInfo};
 
@@ -19,7 +18,7 @@ use tikv_jemallocator::Jemalloc;
 static GLOBAL: Jemalloc = Jemalloc;
 
 /// Aggregate links in the at-mosphere
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 #[command(version, about, long_about = None)]
 struct Args {
     /// Jetstream server to connect to (exclusive with --fixture). Provide either a wss:// URL, or a shorhand value:
@@ -47,9 +46,6 @@ struct Args {
     /// todo: restore this
     #[arg(long, action)]
     pause_rw: bool,
-    /// DEBUG: use an in-memory store instead of fjall
-    #[arg(long, action)]
-    in_mem: bool,
     /// reset the rollup cursor, scrape through missed things in the past (backfill)
     #[arg(long, action)]
     reroll: bool,
@@ -64,56 +60,18 @@ async fn main() -> anyhow::Result<()> {
 
     let args = Args::parse();
     let jetstream = args.jetstream.clone();
-    if args.in_mem {
-        let (read_store, write_store, cursor, sketch_secret) = MemStorage::init(
-            args.data,
-            jetstream,
-            args.jetstream_force,
-            Default::default(),
-        )?;
-        go(
-            args.jetstream,
-            args.jetstream_fixture,
-            args.pause_writer,
-            args.backfill,
-            args.reroll,
-            read_store,
-            write_store,
-            cursor,
-            sketch_secret,
-        )
-        .await?;
-    } else {
-        let (read_store, write_store, cursor, sketch_secret) = FjallStorage::init(
-            args.data,
-            jetstream,
-            args.jetstream_force,
-            Default::default(),
-        )?;
-        go(
-            args.jetstream,
-            args.jetstream_fixture,
-            args.pause_writer,
-            args.backfill,
-            args.reroll,
-            read_store,
-            write_store,
-            cursor,
-            sketch_secret,
-        )
-        .await?;
-    }
-
+    let (read_store, write_store, cursor, sketch_secret) = FjallStorage::init(
+        args.data.clone(),
+        jetstream,
+        args.jetstream_force,
+        Default::default(),
+    )?;
+    go(args, read_store, write_store, cursor, sketch_secret).await?;
     Ok(())
 }
 
-#[allow(clippy::too_many_arguments)]
 async fn go<B: StoreBackground>(
-    jetstream: String,
-    jetstream_fixture: bool,
-    pause_writer: bool,
-    backfill: bool,
-    reroll: bool,
+    args: Args,
     read_store: impl StoreReader + 'static + Clone,
     mut write_store: impl StoreWriter<B> + 'static,
     cursor: Option<Cursor>,
@@ -122,24 +80,26 @@ async fn go<B: StoreBackground>(
     println!("starting server with storage...");
     let serving = server::serve(read_store.clone());
 
-    if pause_writer {
+    if args.pause_writer {
         log::info!("not starting jetstream or the write loop.");
         serving.await.map_err(|e| anyhow::anyhow!(e))?;
         return Ok(());
     }
 
-    let batches = if jetstream_fixture {
-        log::info!("starting with jestream file fixture: {jetstream:?}");
-        file_consumer::consume(jetstream.into(), sketch_secret, cursor).await?
+    let batches = if args.jetstream_fixture {
+        log::info!("starting with jestream file fixture: {:?}", args.jetstream);
+        file_consumer::consume(args.jetstream.into(), sketch_secret, cursor).await?
     } else {
         log::info!(
             "starting consumer with cursor: {cursor:?} from {:?} ago",
             cursor.map(|c| c.elapsed())
         );
-        consumer::consume(&jetstream, cursor, false, sketch_secret).await?
+        consumer::consume(&args.jetstream, cursor, false, sketch_secret).await?
     };
 
-    let rolling = write_store.background_tasks(reroll)?.run(backfill);
+    let rolling = write_store
+        .background_tasks(args.reroll)?
+        .run(args.backfill);
     let storing = write_store.receive_batches(batches);
 
     let stating = do_update_stuff(read_store);

ufos/src/main.rs

@@ -0,0 +1,72 @@
+use crate::Nsid;
+use async_trait::async_trait;
+use dropshot::{
+    ApiEndpointBodyContentType, ExtractorMetadata, HttpError, Query, RequestContext, ServerContext,
+    SharedExtractor,
+};
+use schemars::JsonSchema;
+use serde::Deserialize;
+use std::collections::HashSet;
+
+/// The real type that gets deserialized
+#[derive(Debug, Deserialize, JsonSchema)]
+pub struct MultiCollectionQuery {
+    pub collection: Vec<String>,
+}
+
+/// The fake corresponding type for docs that dropshot won't freak out about a
+/// vec for
+#[derive(Deserialize, JsonSchema)]
+#[allow(dead_code)]
+struct MultiCollectionQueryForDocs {
+    /// One or more collection [NSID](https://atproto.com/specs/nsid)s
+    ///
+    /// Pass this parameter multiple times to specify multiple collections, like
+    /// `collection=app.bsky.feed.like&collection=app.bsky.feed.post`
+    collection: String,
+}
+
+impl TryFrom<MultiCollectionQuery> for HashSet<Nsid> {
+    type Error = HttpError;
+    fn try_from(mcq: MultiCollectionQuery) -> Result<Self, Self::Error> {
+        let mut out = HashSet::with_capacity(mcq.collection.len());
+        for c in mcq.collection {
+            let nsid = Nsid::new(c).map_err(|e| {
+                HttpError::for_bad_request(
+                    None,
+                    format!("failed to convert collection to an NSID: {e:?}"),
+                )
+            })?;
+            out.insert(nsid);
+        }
+        Ok(out)
+    }
+}
+
+// The `SharedExtractor` implementation for Query<QueryType> describes how to
+// construct an instance of `Query<QueryType>` from an HTTP request: namely, by
+// parsing the query string to an instance of `QueryType`.
+#[async_trait]
+impl SharedExtractor for MultiCollectionQuery {
+    async fn from_request<Context: ServerContext>(
+        ctx: &RequestContext<Context>,
+    ) -> Result<MultiCollectionQuery, HttpError> {
+        let raw_query = ctx.request.uri().query().unwrap_or("");
+        let q = serde_qs::from_str(raw_query).map_err(|e| {
+            HttpError::for_bad_request(None, format!("unable to parse query string: {}", e))
+        })?;
+        Ok(q)
+    }
+
+    fn metadata(body_content_type: ApiEndpointBodyContentType) -> ExtractorMetadata {
+        // HACK: query type switcheroo: passing MultiCollectionQuery to
+        // `metadata` would "helpfully" panic because dropshot believes we can
+        // only have scalar types in a  query.
+        //
+        // so instead we have a fake second type whose only job is to look the
+        // same as MultiCollectionQuery exept that it has `String` instead of
+        // `Vec<String>`, which dropshot will accept, and generate ~close-enough
+        // docs for.
+        <Query<MultiCollectionQueryForDocs> as SharedExtractor>::metadata(body_content_type)
+    }
+}

ufos/src/server/collections_query.rs

@@ -0,0 +1,23 @@
+use dropshot::{HttpError, HttpResponseHeaders, HttpResponseOk};
+use schemars::JsonSchema;
+use serde::Serialize;
+
+pub type OkCorsResponse<T> = Result<HttpResponseHeaders<HttpResponseOk<T>>, HttpError>;
+
+/// Helper for constructing Ok responses: return OkCors(T).into()
+/// (not happy with this yet)
+pub struct OkCors<T: Serialize + JsonSchema + Send + Sync>(pub T);
+
+impl<T> From<OkCors<T>> for OkCorsResponse<T>
+where
+    T: Serialize + JsonSchema + Send + Sync,
+{
+    fn from(ok: OkCors<T>) -> OkCorsResponse<T> {
+        let mut res = HttpResponseHeaders::new_unnamed(HttpResponseOk(ok.0));
+        res.headers_mut()
+            .insert("access-control-allow-origin", "*".parse().unwrap());
+        Ok(res)
+    }
+}
+
+// TODO: cors for HttpError