commit 57c0d71d2ed8e74e2354df08acb99bb1fa66707c · nekomimi.pet/microcosm-rs

+1 -1

Makefile

···

       5
        
       	cargo test --all-features

     

       6
        
       

     

       7
        
       fmt:

     

       8
       -
       	cargo fmt --package links --package constellation --package ufos

     

       9
        
       	cargo +nightly fmt --package jetstream

     

       10
        
       

     

       11
        
       clippy:

···

       5
        
       	cargo test --all-features

     

       6
        
       

     

       7
        
       fmt:

     

       8
       +
       	cargo fmt --package links --package constellation --package ufos --package spacedust --package who-am-i

     

       9
        
       	cargo +nightly fmt --package jetstream

     

       10
        
       

     

       11
        
       clippy:

+12 -6

spacedust/src/consumer.rs

···

       1
       -
       use std::sync::Arc;

     

       2
       -
       use tokio_util::sync::CancellationToken;

     

       3
        
       use crate::ClientMessage;

     

       4
        
       use crate::error::ConsumerError;

     

       5
        
       use crate::removable_delay_queue;

     
···

       8
        
           events::{CommitOp, Cursor, EventKind},

     

       9
        
       };

     

       10
        
       use links::collect_links;

     

       0
        
       
     

       11
        
       use tokio::sync::broadcast;

     

       0
        
       
     

       12
        
       

     

       13
        
       const MAX_LINKS_PER_EVENT: usize = 100;

     

       14
        
       

     
···

       61
        
               };

     

       62
        
       

     

       63
        
               // TODO: something a bit more robust

     

       64
       -
               let at_uri = format!("at://{}/{}/{}", &*event.did, &*commit.collection, &*commit.rkey);

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       65
        
       

     

       66
        
               // TODO: keep a buffer and remove quick deletes to debounce notifs

     

       67
        
               // for now we just drop all deletes eek

     

       68
        
               if commit.operation == CommitOp::Delete {

     

       69
       -
                   d.remove_range((at_uri.clone(), 0)..=(at_uri.clone(), MAX_LINKS_PER_EVENT)).await;

     

       0
        
       
     

       70
        
                   continue;

     

       71
        
               }

     

       72
        
               let Some(ref record) = commit.record else {

     
···

       86
        
                   if i >= MAX_LINKS_PER_EVENT {

     

       87
        
                       // todo: indicate if the link limit was reached (-> links omitted)

     

       88
        
                       log::warn!("consumer: event has too many links, ignoring the rest");

     

       89
       -
                       metrics::counter!("consumer_dropped_links", "reason" => "too_many_links").increment(1);

     

       0
        
       
     

       90
        
                       break;

     

       91
        
                   }

     

       92
        
                   let client_message = match ClientMessage::new_link(link, &at_uri, commit) {

     
···

       94
        
                       Err(e) => {

     

       95
        
                           // TODO indicate to clients that a link has been dropped

     

       96
        
                           log::warn!("consumer: failed to serialize link to json: {e:?}");

     

       97
       -
                           metrics::counter!("consumer_dropped_links", "reason" => "failed_to_serialize").increment(1);

     

       0
        
       
     

       98
        
                           continue;

     

       99
        
                       }

     

       100
        
                   };

···

       0
        
       
     

       0
        
       
     

       1
        
       use crate::ClientMessage;

     

       2
        
       use crate::error::ConsumerError;

     

       3
        
       use crate::removable_delay_queue;

     
···

       6
        
           events::{CommitOp, Cursor, EventKind},

     

       7
        
       };

     

       8
        
       use links::collect_links;

     

       9
       +
       use std::sync::Arc;

     

       10
        
       use tokio::sync::broadcast;

     

       11
       +
       use tokio_util::sync::CancellationToken;

     

       12
        
       

     

       13
        
       const MAX_LINKS_PER_EVENT: usize = 100;

     

       14
        
       

     
···

       61
        
               };

     

       62
        
       

     

       63
        
               // TODO: something a bit more robust

     

       64
       +
               let at_uri = format!(

     

       65
       +
                   "at://{}/{}/{}",

     

       66
       +
                   &*event.did, &*commit.collection, &*commit.rkey

     

       67
       +
               );

     

       68
        
       

     

       69
        
               // TODO: keep a buffer and remove quick deletes to debounce notifs

     

       70
        
               // for now we just drop all deletes eek

     

       71
        
               if commit.operation == CommitOp::Delete {

     

       72
       +
                   d.remove_range((at_uri.clone(), 0)..=(at_uri.clone(), MAX_LINKS_PER_EVENT))

     

       73
       +
                       .await;

     

       74
        
                   continue;

     

       75
        
               }

     

       76
        
               let Some(ref record) = commit.record else {

     
···

       90
        
                   if i >= MAX_LINKS_PER_EVENT {

     

       91
        
                       // todo: indicate if the link limit was reached (-> links omitted)

     

       92
        
                       log::warn!("consumer: event has too many links, ignoring the rest");

     

       93
       +
                       metrics::counter!("consumer_dropped_links", "reason" => "too_many_links")

     

       94
       +
                           .increment(1);

     

       95
        
                       break;

     

       96
        
                   }

     

       97
        
                   let client_message = match ClientMessage::new_link(link, &at_uri, commit) {

     
···

       99
        
                       Err(e) => {

     

       100
        
                           // TODO indicate to clients that a link has been dropped

     

       101
        
                           log::warn!("consumer: failed to serialize link to json: {e:?}");

     

       102
       +
                           metrics::counter!("consumer_dropped_links", "reason" => "failed_to_serialize")

     

       103
       +
                               .increment(1);

     

       104
        
                           continue;

     

       105
        
                       }

     

       106
        
                   };

+2 -2

spacedust/src/delay.rs

···

       0
        
       
     

       1
        
       use crate::removable_delay_queue;

     

       2
       -
       use tokio_util::sync::CancellationToken;

     

       3
        
       use tokio::sync::broadcast;

     

       4
       -
       use crate::error::DelayError;

     

       5
        
       

     

       6
        
       pub async fn to_broadcast<T>(

     

       7
        
           source: removable_delay_queue::Output<(String, usize), T>,

···

       1
       +
       use crate::error::DelayError;

     

       2
        
       use crate::removable_delay_queue;

     

       0
        
       
     

       3
        
       use tokio::sync::broadcast;

     

       4
       +
       use tokio_util::sync::CancellationToken;

     

       5
        
       

     

       6
        
       pub async fn to_broadcast<T>(

     

       7
        
           source: removable_delay_queue::Output<(String, usize), T>,

+19 -8

spacedust/src/lib.rs

···

       1
        
       pub mod consumer;

     

       2
        
       pub mod delay;

     

       3
        
       pub mod error;

     

       0
        
       
     

       4
        
       pub mod server;

     

       5
        
       pub mod subscriber;

     

       6
       -
       pub mod removable_delay_queue;

     

       7
        
       

     

       8
       -
       use links::CollectedLink;

     

       9
        
       use jetstream::events::CommitEvent;

     

       10
       -
       use tokio_tungstenite::tungstenite::Message;

     

       11
        
       use serde::{Deserialize, Serialize};

     

       12
        
       use server::MultiSubscribeQuery;

     

       0
        
       
     

       13
        
       

     

       14
        
       #[derive(Debug)]

     

       15
        
       pub struct FilterableProperties {

     
···

       32
        
       }

     

       33
        
       

     

       34
        
       impl ClientMessage {

     

       35
       -
           pub fn new_link(link: CollectedLink, at_uri: &str, commit: &CommitEvent) -> Result<Self, serde_json::Error> {

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       36
        
               let subject_did = link.target.did();

     

       37
        
       

     

       38
        
               let subject = link.target.into_string();

     
···

       61
        
       

     

       62
        
               let message = Message::Text(client_event_json.into());

     

       63
        
       

     

       64
       -
               let properties = FilterableProperties { subject, subject_did, source };

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       65
        
       

     

       66
       -
               Ok(ClientMessage { message, properties })

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       67
        
           }

     

       68
        
       }

     

       69
        
       

     

       70
        
       #[derive(Debug, Serialize)]

     

       71
       -
       #[serde(rename_all="snake_case")]

     

       72
        
       pub struct ClientEvent {

     

       73
       -
           kind: &'static str, // "link"

     

       74
        
           origin: &'static str, // "live", "replay", "backfill"

     

       75
        
           link: ClientLinkEvent,

     

       76
        
       }

···

       1
        
       pub mod consumer;

     

       2
        
       pub mod delay;

     

       3
        
       pub mod error;

     

       4
       +
       pub mod removable_delay_queue;

     

       5
        
       pub mod server;

     

       6
        
       pub mod subscriber;

     

       0
        
       
     

       7
        
       

     

       0
        
       
     

       8
        
       use jetstream::events::CommitEvent;

     

       9
       +
       use links::CollectedLink;

     

       10
        
       use serde::{Deserialize, Serialize};

     

       11
        
       use server::MultiSubscribeQuery;

     

       12
       +
       use tokio_tungstenite::tungstenite::Message;

     

       13
        
       

     

       14
        
       #[derive(Debug)]

     

       15
        
       pub struct FilterableProperties {

     
···

       32
        
       }

     

       33
        
       

     

       34
        
       impl ClientMessage {

     

       35
       +
           pub fn new_link(

     

       36
       +
               link: CollectedLink,

     

       37
       +
               at_uri: &str,

     

       38
       +
               commit: &CommitEvent,

     

       39
       +
           ) -> Result<Self, serde_json::Error> {

     

       40
        
               let subject_did = link.target.did();

     

       41
        
       

     

       42
        
               let subject = link.target.into_string();

     
···

       65
        
       

     

       66
        
               let message = Message::Text(client_event_json.into());

     

       67
        
       

     

       68
       +
               let properties = FilterableProperties {

     

       69
       +
                   subject,

     

       70
       +
                   subject_did,

     

       71
       +
                   source,

     

       72
       +
               };

     

       73
        
       

     

       74
       +
               Ok(ClientMessage {

     

       75
       +
                   message,

     

       76
       +
                   properties,

     

       77
       +
               })

     

       78
        
           }

     

       79
        
       }

     

       80
        
       

     

       81
        
       #[derive(Debug, Serialize)]

     

       82
       +
       #[serde(rename_all = "snake_case")]

     

       83
        
       pub struct ClientEvent {

     

       84
       +
           kind: &'static str,   // "link"

     

       85
        
           origin: &'static str, // "live", "replay", "backfill"

     

       86
        
           link: ClientLinkEvent,

     

       87
        
       }

+11 -6

spacedust/src/main.rs

···

       1
       -
       use spacedust::error::MainTaskError;

     

       2
        
       use spacedust::consumer;

     

       3
       -
       use spacedust::server;

     

       4
        
       use spacedust::delay;

     

       0
        
       
     

       5
        
       use spacedust::removable_delay_queue::removable_delay_queue;

     

       0
        
       
     

       6
        
       

     

       7
        
       use clap::Parser;

     

       8
        
       use metrics_exporter_prometheus::PrometheusBuilder;

     

       0
        
       
     

       9
        
       use tokio::sync::broadcast;

     

       10
        
       use tokio_util::sync::CancellationToken;

     

       11
       -
       use std::time::Duration;

     

       12
        
       

     

       13
        
       /// Aggregate links in the at-mosphere

     

       14
        
       #[derive(Parser, Debug, Clone)]

     
···

       80
        
                   args.jetstream,

     

       81
        
                   None,

     

       82
        
                   args.jetstream_no_zstd,

     

       83
       -
                   consumer_shutdown

     

       84
        
               )

     

       85
       -
                   .await?;

     

       86
        
               Ok(())

     

       87
        
           });

     

       88
        
       

     

       89
        
           let delay_shutdown = shutdown.clone();

     

       90
        
           tasks.spawn(async move {

     

       91
       -
               delay::to_broadcast(delay_queue_receiver, consumer_delayed_sender, delay_shutdown).await?;

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       92
        
               Ok(())

     

       93
        
           });

     

       94

···

       0
        
       
     

       1
        
       use spacedust::consumer;

     

       0
        
       
     

       2
        
       use spacedust::delay;

     

       3
       +
       use spacedust::error::MainTaskError;

     

       4
        
       use spacedust::removable_delay_queue::removable_delay_queue;

     

       5
       +
       use spacedust::server;

     

       6
        
       

     

       7
        
       use clap::Parser;

     

       8
        
       use metrics_exporter_prometheus::PrometheusBuilder;

     

       9
       +
       use std::time::Duration;

     

       10
        
       use tokio::sync::broadcast;

     

       11
        
       use tokio_util::sync::CancellationToken;

     

       0
        
       
     

       12
        
       

     

       13
        
       /// Aggregate links in the at-mosphere

     

       14
        
       #[derive(Parser, Debug, Clone)]

     
···

       80
        
                   args.jetstream,

     

       81
        
                   None,

     

       82
        
                   args.jetstream_no_zstd,

     

       83
       +
                   consumer_shutdown,

     

       84
        
               )

     

       85
       +
               .await?;

     

       86
        
               Ok(())

     

       87
        
           });

     

       88
        
       

     

       89
        
           let delay_shutdown = shutdown.clone();

     

       90
        
           tasks.spawn(async move {

     

       91
       +
               delay::to_broadcast(

     

       92
       +
                   delay_queue_receiver,

     

       93
       +
                   consumer_delayed_sender,

     

       94
       +
                   delay_shutdown,

     

       95
       +
               )

     

       96
       +
               .await?;

     

       97
        
               Ok(())

     

       98
        
           });

     

       99

+15 -11

spacedust/src/removable_delay_queue.rs

···

       1
       -
       use std::ops::RangeBounds;

     

       2
        
       use std::collections::{BTreeMap, VecDeque};

     

       3
       -
       use std::time::{Duration, Instant};

     

       4
       -
       use tokio::sync::Mutex;

     

       5
        
       use std::sync::Arc;

     

       0
        
       
     

       6
        
       use thiserror::Error;

     

       0
        
       
     

       7
        
       

     

       8
        
       #[derive(Debug, Error)]

     

       9
        
       pub enum EnqueueError<T> {

     
···

       17
        
       #[derive(Debug)]

     

       18
        
       struct Queue<K: Key, T> {

     

       19
        
           queue: VecDeque<(Instant, K)>,

     

       20
       -
           items: BTreeMap<K, T>

     

       21
        
       }

     

       22
        
       

     

       23
        
       pub struct Input<K: Key, T> {

     
···

       49
        
           pub async fn remove_range(&self, range: impl RangeBounds<K>) {

     

       50
        
               let n = {

     

       51
        
                   let mut q = self.q.lock().await;

     

       52
       -
                   let keys = q.items.range(range).map(|(k, _)| k).cloned().collect::<Vec<_>>();

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       53
        
                   for k in &keys {

     

       54
        
                       q.items.remove(k);

     

       55
        
                   }

     
···

       94
        
                       } else {

     

       95
        
                           let overshoot = now.saturating_duration_since(expected_release);

     

       96
        
                           metrics::counter!("delay_queue_emit_total", "early" => "no").increment(1);

     

       97
       -
                           metrics::histogram!("delay_queue_emit_overshoot").record(overshoot.as_secs_f64());

     

       0
        
       
     

       98
        
                       }

     

       99
       -
                       return Some(item)

     

       100
        
                   } else if Arc::strong_count(&self.q) == 1 {

     

       101
        
                       return None;

     

       102
        
                   }

     

       103
        
                   // the queue is *empty*, so we need to wait at least as long as the current delay

     

       104
        
                   tokio::time::sleep(self.delay).await;

     

       105
        
                   metrics::counter!("delay_queue_entirely_empty_total").increment(1);

     

       106
       -
               };

     

       107
        
           }

     

       108
        
       }

     

       109
        
       

     

       110
       -
       pub fn removable_delay_queue<K: Key, T>(

     

       111
       -
           delay: Duration,

     

       112
       -
       ) -> (Input<K, T>, Output<K, T>) {

     

       113
        
           let q: Arc<Mutex<Queue<K, T>>> = Arc::new(Mutex::new(Queue {

     

       114
        
               queue: VecDeque::new(),

     

       115
        
               items: BTreeMap::new(),

···

       0
        
       
     

       1
        
       use std::collections::{BTreeMap, VecDeque};

     

       2
       +
       use std::ops::RangeBounds;

     

       0
        
       
     

       3
        
       use std::sync::Arc;

     

       4
       +
       use std::time::{Duration, Instant};

     

       5
        
       use thiserror::Error;

     

       6
       +
       use tokio::sync::Mutex;

     

       7
        
       

     

       8
        
       #[derive(Debug, Error)]

     

       9
        
       pub enum EnqueueError<T> {

     
···

       17
        
       #[derive(Debug)]

     

       18
        
       struct Queue<K: Key, T> {

     

       19
        
           queue: VecDeque<(Instant, K)>,

     

       20
       +
           items: BTreeMap<K, T>,

     

       21
        
       }

     

       22
        
       

     

       23
        
       pub struct Input<K: Key, T> {

     
···

       49
        
           pub async fn remove_range(&self, range: impl RangeBounds<K>) {

     

       50
        
               let n = {

     

       51
        
                   let mut q = self.q.lock().await;

     

       52
       +
                   let keys = q

     

       53
       +
                       .items

     

       54
       +
                       .range(range)

     

       55
       +
                       .map(|(k, _)| k)

     

       56
       +
                       .cloned()

     

       57
       +
                       .collect::<Vec<_>>();

     

       58
        
                   for k in &keys {

     

       59
        
                       q.items.remove(k);

     

       60
        
                   }

     
···

       99
        
                       } else {

     

       100
        
                           let overshoot = now.saturating_duration_since(expected_release);

     

       101
        
                           metrics::counter!("delay_queue_emit_total", "early" => "no").increment(1);

     

       102
       +
                           metrics::histogram!("delay_queue_emit_overshoot")

     

       103
       +
                               .record(overshoot.as_secs_f64());

     

       104
        
                       }

     

       105
       +
                       return Some(item);

     

       106
        
                   } else if Arc::strong_count(&self.q) == 1 {

     

       107
        
                       return None;

     

       108
        
                   }

     

       109
        
                   // the queue is *empty*, so we need to wait at least as long as the current delay

     

       110
        
                   tokio::time::sleep(self.delay).await;

     

       111
        
                   metrics::counter!("delay_queue_entirely_empty_total").increment(1);

     

       112
       +
               }

     

       113
        
           }

     

       114
        
       }

     

       115
        
       

     

       116
       +
       pub fn removable_delay_queue<K: Key, T>(delay: Duration) -> (Input<K, T>, Output<K, T>) {

     

       0
        
       
     

       0
        
       
     

       117
        
           let q: Arc<Mutex<Queue<K, T>>> = Arc::new(Mutex::new(Queue {

     

       118
        
               queue: VecDeque::new(),

     

       119
        
               items: BTreeMap::new(),

+19 -17

spacedust/src/server.rs

···

       0
        
       
     

       1
        
       use crate::error::ServerError;

     

       2
        
       use crate::subscriber::Subscriber;

     

       3
       -
       use metrics::{histogram, counter};

     

       4
       -
       use std::sync::Arc;

     

       5
       -
       use crate::ClientMessage;

     

       0
        
       
     

       0
        
       
     

       6
        
       use http::{

     

       7
       -
           header::{ORIGIN, USER_AGENT},

     

       8
        
           Response, StatusCode,

     

       9
       -
       };

     

       10
       -
       use dropshot::{

     

       11
       -
           Body,

     

       12
       -
           ApiDescription, ConfigDropshot, ConfigLogging, ConfigLoggingLevel, Query, RequestContext,

     

       13
       -
           ServerBuilder, WebsocketConnection, channel, endpoint, HttpResponse,

     

       14
       -
           ApiEndpointBodyContentType, ExtractorMetadata, HttpError, ServerContext,

     

       15
       -
           SharedExtractor,

     

       16
        
       };

     

       0
        
       
     

       0
        
       
     

       17
        
       

     

       0
        
       
     

       18
        
       use schemars::JsonSchema;

     

       19
        
       use serde::{Deserialize, Serialize};

     

       0
        
       
     

       20
        
       use tokio::sync::broadcast;

     

       21
        
       use tokio::time::Instant;

     

       22
        
       use tokio_tungstenite::tungstenite::protocol::{Role, WebSocketConfig};

     

       23
        
       use tokio_util::sync::CancellationToken;

     

       24
       -
       use async_trait::async_trait;

     

       25
       -
       use std::collections::HashSet;

     

       26
        
       

     

       27
        
       const INDEX_HTML: &str = include_str!("../static/index.html");

     

       28
        
       const FAVICON: &[u8] = include_bytes!("../static/favicon.ico");

     
···

       30
        
       pub async fn serve(

     

       31
        
           b: broadcast::Sender<Arc<ClientMessage>>,

     

       32
        
           d: broadcast::Sender<Arc<ClientMessage>>,

     

       33
       -
           shutdown: CancellationToken

     

       34
        
       ) -> Result<(), ServerError> {

     

       35
        
           let config_logging = ConfigLogging::StderrTerminal {

     

       36
        
               level: ConfigLoggingLevel::Info,

     
···

       65
        
           );

     

       66
        
       

     

       67
        
           let sub_shutdown = shutdown.clone();

     

       68
       -
           let ctx = Context { spec, b, d, shutdown: sub_shutdown };

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       69
        
       

     

       70
        
           let server = ServerBuilder::new(api, ctx, log)

     

       71
        
               .config(ConfigDropshot {

     
···

       161
        
       }

     

       162
        
       

     

       163
        
       // TODO: cors for HttpError

     

       164
       -
       

     

       165
        
       

     

       166
        
       /// Serve index page as html

     

       167
        
       #[endpoint {

     
···

       316
        
               upgraded.into_inner(),

     

       317
        
               Role::Server,

     

       318
        
               Some(WebSocketConfig::default().max_message_size(

     

       319
       -
                   Some(10 * 2_usize.pow(20)) // 10MiB, matching jetstream

     

       320
        
               )),

     

       321
        
           )

     

       322
        
           .await;

···

       1
       +
       use crate::ClientMessage;

     

       2
        
       use crate::error::ServerError;

     

       3
        
       use crate::subscriber::Subscriber;

     

       4
       +
       use dropshot::{

     

       5
       +
           ApiDescription, ApiEndpointBodyContentType, Body, ConfigDropshot, ConfigLogging,

     

       6
       +
           ConfigLoggingLevel, ExtractorMetadata, HttpError, HttpResponse, Query, RequestContext,

     

       7
       +
           ServerBuilder, ServerContext, SharedExtractor, WebsocketConnection, channel, endpoint,

     

       8
       +
       };

     

       9
        
       use http::{

     

       0
        
       
     

       10
        
           Response, StatusCode,

     

       11
       +
           header::{ORIGIN, USER_AGENT},

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       12
        
       };

     

       13
       +
       use metrics::{counter, histogram};

     

       14
       +
       use std::sync::Arc;

     

       15
        
       

     

       16
       +
       use async_trait::async_trait;

     

       17
        
       use schemars::JsonSchema;

     

       18
        
       use serde::{Deserialize, Serialize};

     

       19
       +
       use std::collections::HashSet;

     

       20
        
       use tokio::sync::broadcast;

     

       21
        
       use tokio::time::Instant;

     

       22
        
       use tokio_tungstenite::tungstenite::protocol::{Role, WebSocketConfig};

     

       23
        
       use tokio_util::sync::CancellationToken;

     

       0
        
       
     

       0
        
       
     

       24
        
       

     

       25
        
       const INDEX_HTML: &str = include_str!("../static/index.html");

     

       26
        
       const FAVICON: &[u8] = include_bytes!("../static/favicon.ico");

     
···

       28
        
       pub async fn serve(

     

       29
        
           b: broadcast::Sender<Arc<ClientMessage>>,

     

       30
        
           d: broadcast::Sender<Arc<ClientMessage>>,

     

       31
       +
           shutdown: CancellationToken,

     

       32
        
       ) -> Result<(), ServerError> {

     

       33
        
           let config_logging = ConfigLogging::StderrTerminal {

     

       34
        
               level: ConfigLoggingLevel::Info,

     
···

       63
        
           );

     

       64
        
       

     

       65
        
           let sub_shutdown = shutdown.clone();

     

       66
       +
           let ctx = Context {

     

       67
       +
               spec,

     

       68
       +
               b,

     

       69
       +
               d,

     

       70
       +
               shutdown: sub_shutdown,

     

       71
       +
           };

     

       72
        
       

     

       73
        
           let server = ServerBuilder::new(api, ctx, log)

     

       74
        
               .config(ConfigDropshot {

     
···

       164
        
       }

     

       165
        
       

     

       166
        
       // TODO: cors for HttpError

     

       0
        
       
     

       167
        
       

     

       168
        
       /// Serve index page as html

     

       169
        
       #[endpoint {

     
···

       318
        
               upgraded.into_inner(),

     

       319
        
               Role::Server,

     

       320
        
               Some(WebSocketConfig::default().max_message_size(

     

       321
       +
                   Some(10 * 2_usize.pow(20)), // 10MiB, matching jetstream

     

       322
        
               )),

     

       323
        
           )

     

       324
        
           .await;

+23 -26

spacedust/src/subscriber.rs

···

       1
        
       use crate::error::SubscriberUpdateError;

     

       2
       -
       use std::sync::Arc;

     

       3
       -
       use tokio::time::interval;

     

       4
       -
       use std::time::Duration;

     

       5
       -
       use futures::StreamExt;

     

       6
        
       use crate::{ClientMessage, FilterableProperties, SubscriberSourcedMessage};

     

       7
       -
       use crate::server::MultiSubscribeQuery;

     

       8
        
       use futures::SinkExt;

     

       0
        
       
     

       9
        
       use std::error::Error;

     

       0
        
       
     

       0
        
       
     

       10
        
       use tokio::sync::broadcast::{self, error::RecvError};

     

       0
        
       
     

       11
        
       use tokio_tungstenite::{WebSocketStream, tungstenite::Message};

     

       12
        
       use tokio_util::sync::CancellationToken;

     

       13
       -
       use dropshot::WebsocketConnectionRaw;

     

       14
        
       

     

       15
        
       const PING_PERIOD: Duration = Duration::from_secs(30);

     

       16
        
       

     
···

       20
        
       }

     

       21
        
       

     

       22
        
       impl Subscriber {

     

       23
       -
           pub fn new(

     

       24
       -
               query: MultiSubscribeQuery,

     

       25
       -
               shutdown: CancellationToken,

     

       26
       -
           ) -> Self {

     

       27
        
               Self { query, shutdown }

     

       28
        
           }

     

       29
        
       

     

       30
        
           pub async fn start(

     

       31
        
               mut self,

     

       32
        
               ws: WebSocketStream<WebsocketConnectionRaw>,

     

       33
       -
               mut receiver: broadcast::Receiver<Arc<ClientMessage>>

     

       34
        
           ) -> Result<(), Box<dyn Error>> {

     

       35
        
               let mut ping_state = None;

     

       36
        
               let (mut ws_sender, mut ws_receiver) = ws.split();

     
···

       83
        
                                   // TODO: send client an explanation

     

       84
        
                                   self.shutdown.cancel();

     

       85
        
                               }

     

       0
        
       
     

       86
        
                           },

     

       87
        
                           Some(Ok(m)) => log::trace!("subscriber sent an unexpected message: {m:?}"),

     

       88
        
                           Some(Err(e)) => {

     
···

       122
        
               Ok(())

     

       123
        
           }

     

       124
        
       

     

       125
       -
           fn filter(

     

       126
       -
               &self,

     

       127
       -
               properties: &FilterableProperties,

     

       128
       -
           ) -> bool {

     

       129
        
               let query = &self.query;

     

       130
        
       

     

       131
        
               // subject + subject DIDs are logical OR

     

       132
       -
               if !(

     

       133
       -
                   query.wanted_subjects.is_empty() && query.wanted_subject_dids.is_empty() ||

     

       134
       -
                   query.wanted_subjects.contains(&properties.subject) ||

     

       135
       -
                   properties.subject_did.as_ref().map(|did| query.wanted_subject_dids.contains(did)).unwrap_or(false)

     

       136
       -
               ) { // wowwww ^^ fix that

     

       137
       -
                   return false

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       138
        
               }

     

       139
        
       

     

       140
        
               // subjects together with sources are logical AND

     

       141
        
               if !(query.wanted_sources.is_empty() || query.wanted_sources.contains(&properties.source)) {

     

       142
       -
                   return false

     

       143
        
               }

     

       144
        
       

     

       145
        
               true

     

       146
        
           }

     

       147
        
       }

     

       148
        
       

     

       149
       -
       

     

       150
       -
       

     

       151
        
       impl MultiSubscribeQuery {

     

       152
        
           pub fn update_from_raw(&mut self, s: &str) -> Result<(), SubscriberUpdateError> {

     

       153
       -
               let SubscriberSourcedMessage::OptionsUpdate(opts) = serde_json::from_str(s)

     

       154
       -
                   .map_err(SubscriberUpdateError::FailedToParseMessage)?;

     

       155
        
               if opts.wanted_sources.len() > 1_000 {

     

       156
        
                   return Err(SubscriberUpdateError::TooManySourcesWanted);

     

       157
        
               }

···

       1
        
       use crate::error::SubscriberUpdateError;

     

       2
       +
       use crate::server::MultiSubscribeQuery;

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       3
        
       use crate::{ClientMessage, FilterableProperties, SubscriberSourcedMessage};

     

       4
       +
       use dropshot::WebsocketConnectionRaw;

     

       5
        
       use futures::SinkExt;

     

       6
       +
       use futures::StreamExt;

     

       7
        
       use std::error::Error;

     

       8
       +
       use std::sync::Arc;

     

       9
       +
       use std::time::Duration;

     

       10
        
       use tokio::sync::broadcast::{self, error::RecvError};

     

       11
       +
       use tokio::time::interval;

     

       12
        
       use tokio_tungstenite::{WebSocketStream, tungstenite::Message};

     

       13
        
       use tokio_util::sync::CancellationToken;

     

       0
        
       
     

       14
        
       

     

       15
        
       const PING_PERIOD: Duration = Duration::from_secs(30);

     

       16
        
       

     
···

       20
        
       }

     

       21
        
       

     

       22
        
       impl Subscriber {

     

       23
       +
           pub fn new(query: MultiSubscribeQuery, shutdown: CancellationToken) -> Self {

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       24
        
               Self { query, shutdown }

     

       25
        
           }

     

       26
        
       

     

       27
        
           pub async fn start(

     

       28
        
               mut self,

     

       29
        
               ws: WebSocketStream<WebsocketConnectionRaw>,

     

       30
       +
               mut receiver: broadcast::Receiver<Arc<ClientMessage>>,

     

       31
        
           ) -> Result<(), Box<dyn Error>> {

     

       32
        
               let mut ping_state = None;

     

       33
        
               let (mut ws_sender, mut ws_receiver) = ws.split();

     
···

       80
        
                                   // TODO: send client an explanation

     

       81
        
                                   self.shutdown.cancel();

     

       82
        
                               }

     

       83
       +
                               log::trace!("subscriber updated with opts: {:?}", self.query);

     

       84
        
                           },

     

       85
        
                           Some(Ok(m)) => log::trace!("subscriber sent an unexpected message: {m:?}"),

     

       86
        
                           Some(Err(e)) => {

     
···

       120
        
               Ok(())

     

       121
        
           }

     

       122
        
       

     

       123
       +
           fn filter(&self, properties: &FilterableProperties) -> bool {

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       124
        
               let query = &self.query;

     

       125
        
       

     

       126
        
               // subject + subject DIDs are logical OR

     

       127
       +
               if !(query.wanted_subjects.is_empty() && query.wanted_subject_dids.is_empty()

     

       128
       +
                   || query.wanted_subjects.contains(&properties.subject)

     

       129
       +
                   || properties

     

       130
       +
                       .subject_did

     

       131
       +
                       .as_ref()

     

       132
       +
                       .map(|did| query.wanted_subject_dids.contains(did))

     

       133
       +
                       .unwrap_or(false))

     

       134
       +
               {

     

       135
       +
                   // wowwww ^^ fix that

     

       136
       +
                   return false;

     

       137
        
               }

     

       138
        
       

     

       139
        
               // subjects together with sources are logical AND

     

       140
        
               if !(query.wanted_sources.is_empty() || query.wanted_sources.contains(&properties.source)) {

     

       141
       +
                   return false;

     

       142
        
               }

     

       143
        
       

     

       144
        
               true

     

       145
        
           }

     

       146
        
       }

     

       147
        
       

     

       0
        
       
     

       0
        
       
     

       148
        
       impl MultiSubscribeQuery {

     

       149
        
           pub fn update_from_raw(&mut self, s: &str) -> Result<(), SubscriberUpdateError> {

     

       150
       +
               let SubscriberSourcedMessage::OptionsUpdate(opts) =

     

       151
       +
                   serde_json::from_str(s).map_err(SubscriberUpdateError::FailedToParseMessage)?;

     

       152
        
               if opts.wanted_sources.len() > 1_000 {

     

       153
        
                   return Err(SubscriberUpdateError::TooManySourcesWanted);

     

       154
        
               }

+1 -1

who-am-i/Cargo.toml

···

       4
        
       edition = "2024"

     

       5
        
       

     

       6
        
       [dependencies]

     

       7
       -
       atrium-api = { version = "0.25.4", default-features = false, features = ["tokio", "agent"] }

     

       8
        
       atrium-identity = "0.1.5"

     

       9
        
       atrium-oauth = "0.1.3"

     

       10
        
       clap = { version = "4.5.40", features = ["derive"] }

···

       4
        
       edition = "2024"

     

       5
        
       

     

       6
        
       [dependencies]

     

       7
       +
       atrium-api = { version = "0.25.4", default-features = false }

     

       8
        
       atrium-identity = "0.1.5"

     

       9
        
       atrium-oauth = "0.1.3"

     

       10
        
       clap = { version = "4.5.40", features = ["derive"] }

+1 -1

who-am-i/src/lib.rs

···

       3
        
       mod server;

     

       4
        
       

     

       5
        
       pub use dns_resolver::HickoryDnsTxtResolver;

     

       0
        
       
     

       6
        
       pub use server::serve;

     

       7
       -
       pub use oauth::{Client, client, authorize};

···

       3
        
       mod server;

     

       4
        
       

     

       5
        
       pub use dns_resolver::HickoryDnsTxtResolver;

     

       6
       +
       pub use oauth::{Client, authorize, client};

     

       7
        
       pub use server::serve;

     

       0

+1 -1

who-am-i/src/main.rs

···

       1
       -
       use who_am_i::serve;

     

       2
        
       use tokio_util::sync::CancellationToken;

     

       0
        
       
     

       3
        
       

     

       4
        
       #[tokio::main]

     

       5
        
       async fn main() {

···

       0
        
       
     

       1
        
       use tokio_util::sync::CancellationToken;

     

       2
       +
       use who_am_i::serve;

     

       3
        
       

     

       4
        
       #[tokio::main]

     

       5
        
       async fn main() {

+14 -17

who-am-i/src/oauth.rs

···

       0
        
       
     

       1
        
       use atrium_identity::{

     

       2
        
           did::{CommonDidResolver, CommonDidResolverConfig, DEFAULT_PLC_DIRECTORY_URL},

     

       3
        
           handle::{AtprotoHandleResolver, AtprotoHandleResolverConfig},

     

       4
        
       };

     

       5
        
       use atrium_oauth::{

     

       6
       -
           AuthorizeOptions,

     

       0
        
       
     

       7
        
           store::{session::MemorySessionStore, state::MemoryStateStore},

     

       8
       -
           AtprotoLocalhostClientMetadata, DefaultHttpClient, KnownScope, OAuthClient, OAuthClientConfig,

     

       9
       -
           OAuthResolverConfig, Scope,

     

       10
        
       };

     

       11
        
       use std::sync::Arc;

     

       12
       -
       use crate::HickoryDnsTxtResolver;

     

       13
        
       

     

       14
        
       pub type Client = OAuthClient<

     

       15
        
           MemoryStateStore,

     
···

       23
        
           let config = OAuthClientConfig {

     

       24
        
               client_metadata: AtprotoLocalhostClientMetadata {

     

       25
        
                   redirect_uris: Some(vec![String::from("http://127.0.0.1:9997/authorized")]),

     

       26
       -
                   scopes: Some(vec![

     

       27
       -
                       Scope::Known(KnownScope::Atproto),

     

       28
       -
                   ]),

     

       29
        
               },

     

       30
        
               keys: None,

     

       31
        
               resolver: OAuthResolverConfig {

     
···

       52
        
       }

     

       53
        
       

     

       54
        
       pub async fn authorize(client: &Client, handle: &str) -> String {

     

       55
       -
           let Ok(url) = client.authorize(

     

       56
       -
               handle,

     

       57
       -
               AuthorizeOptions {

     

       58
       -
                   scopes: vec![

     

       59
       -
                       Scope::Known(KnownScope::Atproto),

     

       60
       -
                   ],

     

       61
       -
                   ..Default::default()

     

       62
       -
               },

     

       63
       -
           )

     

       64
       -
           .await else {

     

       65
        
               panic!("failed to authorize");

     

       66
        
           };

     

       67
        
           url

···

       1
       +
       use crate::HickoryDnsTxtResolver;

     

       2
        
       use atrium_identity::{

     

       3
        
           did::{CommonDidResolver, CommonDidResolverConfig, DEFAULT_PLC_DIRECTORY_URL},

     

       4
        
           handle::{AtprotoHandleResolver, AtprotoHandleResolverConfig},

     

       5
        
       };

     

       6
        
       use atrium_oauth::{

     

       7
       +
           AtprotoLocalhostClientMetadata, AuthorizeOptions, DefaultHttpClient, KnownScope, OAuthClient,

     

       8
       +
           OAuthClientConfig, OAuthResolverConfig, Scope,

     

       9
        
           store::{session::MemorySessionStore, state::MemoryStateStore},

     

       0
        
       
     

       0
        
       
     

       10
        
       };

     

       11
        
       use std::sync::Arc;

     

       0
        
       
     

       12
        
       

     

       13
        
       pub type Client = OAuthClient<

     

       14
        
           MemoryStateStore,

     
···

       22
        
           let config = OAuthClientConfig {

     

       23
        
               client_metadata: AtprotoLocalhostClientMetadata {

     

       24
        
                   redirect_uris: Some(vec![String::from("http://127.0.0.1:9997/authorized")]),

     

       25
       +
                   scopes: Some(vec![Scope::Known(KnownScope::Atproto)]),

     

       0
        
       
     

       0
        
       
     

       26
        
               },

     

       27
        
               keys: None,

     

       28
        
               resolver: OAuthResolverConfig {

     
···

       49
        
       }

     

       50
        
       

     

       51
        
       pub async fn authorize(client: &Client, handle: &str) -> String {

     

       52
       +
           let Ok(url) = client

     

       53
       +
               .authorize(

     

       54
       +
                   handle,

     

       55
       +
                   AuthorizeOptions {

     

       56
       +
                       scopes: vec![Scope::Known(KnownScope::Atproto)],

     

       57
       +
                       ..Default::default()

     

       58
       +
                   },

     

       59
       +
               )

     

       60
       +
               .await

     

       61
       +
           else {

     

       62
        
               panic!("failed to authorize");

     

       63
        
           };

     

       64
        
           url

+16 -18

who-am-i/src/server.rs

···

       1
       -
       

     

       2
        
       use atrium_api::agent::SessionManager;

     

       3
       -
       use std::error::Error;

     

       4
       -
       use metrics::{histogram, counter};

     

       5
       -
       use std::sync::Arc;

     

       0
        
       
     

       0
        
       
     

       6
        
       use http::{

     

       7
       -
           header::{ORIGIN, USER_AGENT},

     

       8
        
           Response, StatusCode,

     

       0
        
       
     

       9
        
       };

     

       10
       -
       use dropshot::{

     

       11
       -
           Body, HttpResponseSeeOther, http_response_see_other,

     

       12
       -
           ApiDescription, ConfigDropshot, ConfigLogging, ConfigLoggingLevel, RequestContext,

     

       13
       -
           ServerBuilder, endpoint, HttpResponse, HttpError, ServerContext, Query,

     

       14
       -
       };

     

       15
        
       

     

       16
        
       use atrium_oauth::CallbackParams;

     

       17
        
       use schemars::JsonSchema;

     
···

       19
        
       use tokio::time::Instant;

     

       20
        
       use tokio_util::sync::CancellationToken;

     

       21
        
       

     

       22
       -
       use crate::{Client, client, authorize};

     

       23
        
       

     

       24
        
       const INDEX_HTML: &str = include_str!("../static/index.html");

     

       25
        
       const FAVICON: &[u8] = include_bytes!("../static/favicon.ico");

     

       26
        
       

     

       27
       -
       pub async fn serve(

     

       28
       -
           shutdown: CancellationToken

     

       29
       -
       ) -> Result<(), Box<dyn Error + Send + Sync>> {

     

       30
        
           let config_logging = ConfigLogging::StderrTerminal {

     

       31
        
               level: ConfigLoggingLevel::Info,

     

       32
        
           };

     

       33
        
       

     

       34
       -
           let log = config_logging

     

       35
       -
               .to_logger("example-basic")?;

     

       36
        
       

     

       37
        
           let mut api = ApiDescription::new();

     

       38
        
           api.register(index).unwrap();

     
···

       58
        
               .json()?,

     

       59
        
           );

     

       60
        
       

     

       61
       -
           let ctx = Context { spec, client: client().into() };

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       62
        
       

     

       63
        
           let server = ServerBuilder::new(api, ctx, log)

     

       64
        
               .config(ConfigDropshot {

     
···

       152
        
       }

     

       153
        
       

     

       154
        
       // TODO: cors for HttpError

     

       155
       -
       

     

       156
        
       

     

       157
        
       /// Serve index page as html

     

       158
        
       #[endpoint {

···

       0
        
       
     

       1
        
       use atrium_api::agent::SessionManager;

     

       2
       +
       use dropshot::{

     

       3
       +
           ApiDescription, Body, ConfigDropshot, ConfigLogging, ConfigLoggingLevel, HttpError,

     

       4
       +
           HttpResponse, HttpResponseSeeOther, Query, RequestContext, ServerBuilder, ServerContext,

     

       5
       +
           endpoint, http_response_see_other,

     

       6
       +
       };

     

       7
        
       use http::{

     

       0
        
       
     

       8
        
           Response, StatusCode,

     

       9
       +
           header::{ORIGIN, USER_AGENT},

     

       10
        
       };

     

       11
       +
       use metrics::{counter, histogram};

     

       12
       +
       use std::error::Error;

     

       13
       +
       use std::sync::Arc;

     

       0
        
       
     

       0
        
       
     

       14
        
       

     

       15
        
       use atrium_oauth::CallbackParams;

     

       16
        
       use schemars::JsonSchema;

     
···

       18
        
       use tokio::time::Instant;

     

       19
        
       use tokio_util::sync::CancellationToken;

     

       20
        
       

     

       21
       +
       use crate::{Client, authorize, client};

     

       22
        
       

     

       23
        
       const INDEX_HTML: &str = include_str!("../static/index.html");

     

       24
        
       const FAVICON: &[u8] = include_bytes!("../static/favicon.ico");

     

       25
        
       

     

       26
       +
       pub async fn serve(shutdown: CancellationToken) -> Result<(), Box<dyn Error + Send + Sync>> {

     

       0
        
       
     

       0
        
       
     

       27
        
           let config_logging = ConfigLogging::StderrTerminal {

     

       28
        
               level: ConfigLoggingLevel::Info,

     

       29
        
           };

     

       30
        
       

     

       31
       +
           let log = config_logging.to_logger("example-basic")?;

     

       0
        
       
     

       32
        
       

     

       33
        
           let mut api = ApiDescription::new();

     

       34
        
           api.register(index).unwrap();

     
···

       54
        
               .json()?,

     

       55
        
           );

     

       56
        
       

     

       57
       +
           let ctx = Context {

     

       58
       +
               spec,

     

       59
       +
               client: client().into(),

     

       60
       +
           };

     

       61
        
       

     

       62
        
           let server = ServerBuilder::new(api, ctx, log)

     

       63
        
               .config(ConfigDropshot {

     
···

       151
        
       }

     

       152
        
       

     

       153
        
       // TODO: cors for HttpError

     

       0
        
       
     

       154
        
       

     

       155
        
       /// Serve index page as html

     

       156
        
       #[endpoint {