commit 9d924ff7036cd6f3a80976fbfe8ed5dcc9eaac83 · nekomimi.pet/microcosm-rs

+10 -4

Cargo.lock

···

       839
        
       

     

       840
        
       [[package]]

     

       841
        
       name = "ctrlc"

     

       842
       -
       version = "3.4.6"

     

       843
        
       source = "registry+https://github.com/rust-lang/crates.io-index"

     

       844
       -
       checksum = "697b5419f348fd5ae2478e8018cb016c00a5881c7f46c717de98ffd135a5651c"

     

       845
        
       dependencies = [

     

       846
        
        "nix",

     

       847
        
        "windows-sys 0.59.0",

     
···

       2374
        
       

     

       2375
        
       [[package]]

     

       2376
        
       name = "nix"

     

       2377
       -
       version = "0.29.0"

     

       2378
        
       source = "registry+https://github.com/rust-lang/crates.io-index"

     

       2379
       -
       checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46"

     

       2380
        
       dependencies = [

     

       2381
        
        "bitflags",

     

       2382
        
        "cfg-if",

     
···

       3383
        
       dependencies = [

     

       3384
        
        "async-trait",

     

       3385
        
        "clap",

     

       0
        
       
     

       3386
        
        "dropshot",

     

       0
        
       
     

       3387
        
        "futures",

     

       3388
        
        "http",

     

       3389
        
        "jetstream",

     

       3390
        
        "links",

     

       0
        
       
     

       3391
        
        "metrics",

     

       0
        
       
     

       3392
        
        "schemars",

     

       3393
        
        "semver",

     

       3394
        
        "serde",

     

       3395
        
        "serde_json",

     

       3396
        
        "serde_qs",

     

       0
        
       
     

       3397
        
        "tinyjson",

     

       3398
        
        "tokio",

     

       3399
        
        "tokio-tungstenite 0.27.0",

     

       0
        
       
     

       3400
        
       ]

     

       3401
        
       

     

       3402
        
       [[package]]

···

       839
        
       

     

       840
        
       [[package]]

     

       841
        
       name = "ctrlc"

     

       842
       +
       version = "3.4.7"

     

       843
        
       source = "registry+https://github.com/rust-lang/crates.io-index"

     

       844
       +
       checksum = "46f93780a459b7d656ef7f071fe699c4d3d2cb201c4b24d085b6ddc505276e73"

     

       845
        
       dependencies = [

     

       846
        
        "nix",

     

       847
        
        "windows-sys 0.59.0",

     
···

       2374
        
       

     

       2375
        
       [[package]]

     

       2376
        
       name = "nix"

     

       2377
       +
       version = "0.30.1"

     

       2378
        
       source = "registry+https://github.com/rust-lang/crates.io-index"

     

       2379
       +
       checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6"

     

       2380
        
       dependencies = [

     

       2381
        
        "bitflags",

     

       2382
        
        "cfg-if",

     
···

       3383
        
       dependencies = [

     

       3384
        
        "async-trait",

     

       3385
        
        "clap",

     

       3386
       +
        "ctrlc",

     

       3387
        
        "dropshot",

     

       3388
       +
        "env_logger",

     

       3389
        
        "futures",

     

       3390
        
        "http",

     

       3391
        
        "jetstream",

     

       3392
        
        "links",

     

       3393
       +
        "log",

     

       3394
        
        "metrics",

     

       3395
       +
        "rand 0.9.1",

     

       3396
        
        "schemars",

     

       3397
        
        "semver",

     

       3398
        
        "serde",

     

       3399
        
        "serde_json",

     

       3400
        
        "serde_qs",

     

       3401
       +
        "thiserror 2.0.12",

     

       3402
        
        "tinyjson",

     

       3403
        
        "tokio",

     

       3404
        
        "tokio-tungstenite 0.27.0",

     

       3405
       +
        "tokio-util",

     

       3406
        
       ]

     

       3407
        
       

     

       3408
        
       [[package]]

spacedust/Cargo.toml

···

       6
        
       [dependencies]

     

       7
        
       async-trait = "0.1.88"

     

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

     

       0
        
       
     

       9
        
       dropshot = "0.16.2"

     

       0
        
       
     

       10
        
       futures = "0.3.31"

     

       11
        
       http = "1.3.1"

     

       12
        
       jetstream = { path = "../jetstream", features = ["metrics"] }

     

       13
        
       links = { path = "../links" }

     

       0
        
       
     

       14
        
       metrics = "0.24.2"

     

       0
        
       
     

       15
        
       schemars = "0.8.22"

     

       16
        
       semver = "1.0.26"

     

       17
        
       serde = { version = "1.0.219", features = ["derive"] }

     

       18
        
       serde_json = "1.0.140"

     

       19
        
       serde_qs = "1.0.0-rc.3"

     

       0
        
       
     

       20
        
       tinyjson = "2.5.1"

     

       21
        
       tokio = { version = "1.45.1", features = ["full"] }

     

       22
        
       tokio-tungstenite = "0.27.0"

     

       0

···

       6
        
       [dependencies]

     

       7
        
       async-trait = "0.1.88"

     

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

     

       9
       +
       ctrlc = "3.4.7"

     

       10
        
       dropshot = "0.16.2"

     

       11
       +
       env_logger = "0.11.8"

     

       12
        
       futures = "0.3.31"

     

       13
        
       http = "1.3.1"

     

       14
        
       jetstream = { path = "../jetstream", features = ["metrics"] }

     

       15
        
       links = { path = "../links" }

     

       16
       +
       log = "0.4.27"

     

       17
        
       metrics = "0.24.2"

     

       18
       +
       rand = "0.9.1"

     

       19
        
       schemars = "0.8.22"

     

       20
        
       semver = "1.0.26"

     

       21
        
       serde = { version = "1.0.219", features = ["derive"] }

     

       22
        
       serde_json = "1.0.140"

     

       23
        
       serde_qs = "1.0.0-rc.3"

     

       24
       +
       thiserror = "2.0.12"

     

       25
        
       tinyjson = "2.5.1"

     

       26
        
       tokio = { version = "1.45.1", features = ["full"] }

     

       27
        
       tokio-tungstenite = "0.27.0"

     

       28
       +
       tokio-util = "0.7.15"

+24 -11

spacedust/src/consumer.rs

···

       0
        
       
     

       1
        
       use crate::LinkEvent;

     

       0
        
       
     

       2
        
       use jetstream::{

     

       3
        
           DefaultJetstreamEndpoints, JetstreamCompression, JetstreamConfig, JetstreamConnector,

     

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

     

       5
        
       };

     

       6
        
       use links::collect_links;

     

       7
       -
       use std::error::Error;

     

       8
        
       use tokio::sync::broadcast;

     

       9
        
       

     

       10
        
       const MAX_LINKS_PER_EVENT: usize = 100;

     

       11
        
       

     

       12
        
       pub async fn consume(

     

       13
        
           b: broadcast::Sender<LinkEvent>,

     

       14
       -
           jetstream_endpoint: &str,

     

       15
        
           cursor: Option<Cursor>,

     

       16
        
           no_zstd: bool,

     

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

     

       18
       -
           let endpoint = DefaultJetstreamEndpoints::endpoint_or_shortcut(jetstream_endpoint);

     

       0
        
       
     

       19
        
           if endpoint == jetstream_endpoint {

     

       20
       -
               std::println!("connecting to jetstream at {endpoint}");

     

       21
        
           } else {

     

       22
       -
               std::println!("connecting to jetstream at {jetstream_endpoint} => {endpoint}");

     

       23
        
           }

     

       24
        
           let config: JetstreamConfig = JetstreamConfig {

     

       25
        
               endpoint,

     
···

       36
        
               .connect_cursor(cursor)

     

       37
        
               .await?;

     

       38
        
       

     

       39
       -
           while let Some(event) = receiver.recv().await {

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       40
        
               if event.kind != EventKind::Commit {

     

       41
        
                   continue;

     

       42
        
               }

     

       43
        
               let Some(commit) = event.commit else {

     

       44
       -
                   eprintln!("jetstream commit event missing commit data, ignoring");

     

       45
        
                   continue;

     

       46
        
               };

     

       47
        
       

     
···

       51
        
                   continue;

     

       52
        
               }

     

       53
        
               let Some(record) = commit.record else {

     

       54
       -
                   eprintln!("jetstream commit update/delete missing record, ignoring");

     

       55
        
                   continue;

     

       56
        
               };

     

       57
        
       

     
···

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

     

       61
        
               for (i, link) in collect_links(&jv).into_iter().enumerate() {

     

       62
        
                   if i >= MAX_LINKS_PER_EVENT {

     

       63
       -
                       eprintln!("jetstream event has too many links, ignoring the rest");

     

       64
        
                       break;

     

       65
        
                   }

     

       66
        
                   let link_ev = LinkEvent {

     
···

       79
        
               }

     

       80
        
           }

     

       81
        
       

     

       82
       -
           Err("jetstream consumer ended".into())

     

       83
        
       }

···

       1
       +
       use tokio_util::sync::CancellationToken;

     

       2
        
       use crate::LinkEvent;

     

       3
       +
       use crate::error::ConsumerError;

     

       4
        
       use jetstream::{

     

       5
        
           DefaultJetstreamEndpoints, JetstreamCompression, JetstreamConfig, JetstreamConnector,

     

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

     

       7
        
       };

     

       8
        
       use links::collect_links;

     

       0
        
       
     

       9
        
       use tokio::sync::broadcast;

     

       10
        
       

     

       11
        
       const MAX_LINKS_PER_EVENT: usize = 100;

     

       12
        
       

     

       13
        
       pub async fn consume(

     

       14
        
           b: broadcast::Sender<LinkEvent>,

     

       15
       +
           jetstream_endpoint: String,

     

       16
        
           cursor: Option<Cursor>,

     

       17
        
           no_zstd: bool,

     

       18
       +
           shutdown: CancellationToken,

     

       19
       +
       ) -> Result<(), ConsumerError> {

     

       20
       +
           let endpoint = DefaultJetstreamEndpoints::endpoint_or_shortcut(&jetstream_endpoint);

     

       21
        
           if endpoint == jetstream_endpoint {

     

       22
       +
               log::info!("connecting to jetstream at {endpoint}");

     

       23
        
           } else {

     

       24
       +
               log::info!("connecting to jetstream at {jetstream_endpoint} => {endpoint}");

     

       25
        
           }

     

       26
        
           let config: JetstreamConfig = JetstreamConfig {

     

       27
        
               endpoint,

     
···

       38
        
               .connect_cursor(cursor)

     

       39
        
               .await?;

     

       40
        
       

     

       41
       +
           log::info!("receiving jetstream messages..");

     

       42
       +
           loop {

     

       43
       +
               if shutdown.is_cancelled() {

     

       44
       +
                   log::info!("exiting consumer for shutdown");

     

       45
       +
                   break;

     

       46
       +
               }

     

       47
       +
               let Some(event) = receiver.recv().await else {

     

       48
       +
                   log::error!("could not receive jetstream event, shutting down...");

     

       49
       +
                   shutdown.cancel();

     

       50
       +
                   break;

     

       51
       +
               };

     

       52
       +
       

     

       53
        
               if event.kind != EventKind::Commit {

     

       54
        
                   continue;

     

       55
        
               }

     

       56
        
               let Some(commit) = event.commit else {

     

       57
       +
                   log::warn!("jetstream commit event missing commit data, ignoring");

     

       58
        
                   continue;

     

       59
        
               };

     

       60
        
       

     
···

       64
        
                   continue;

     

       65
        
               }

     

       66
        
               let Some(record) = commit.record else {

     

       67
       +
                   log::warn!("jetstream commit update/delete missing record, ignoring");

     

       68
        
                   continue;

     

       69
        
               };

     

       70
        
       

     
···

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

     

       74
        
               for (i, link) in collect_links(&jv).into_iter().enumerate() {

     

       75
        
                   if i >= MAX_LINKS_PER_EVENT {

     

       76
       +
                       log::warn!("jetstream event has too many links, ignoring the rest");

     

       77
        
                       break;

     

       78
        
                   }

     

       79
        
                   let link_ev = LinkEvent {

     
···

       92
        
               }

     

       93
        
           }

     

       94
        
       

     

       95
       +
           Err(ConsumerError::JetstreamEnded)

     

       96
        
       }

+13

spacedust/src/error.rs

···

       1
       +
       use thiserror::Error;

     

       2
       +
       

     

       3
       +
       #[derive(Debug, Error)]

     

       4
       +
       pub enum ConsumerError {

     

       5
       +
           #[error(transparent)]

     

       6
       +
           JetstreamConnectionError(#[from] jetstream::error::ConnectionError),

     

       7
       +
           #[error(transparent)]

     

       8
       +
           JetstreamConfigValidationError(#[from] jetstream::error::ConfigValidationError),

     

       9
       +
           #[error(transparent)]

     

       10
       +
           JsonParseError(#[from] tinyjson::JsonParseError),

     

       11
       +
           #[error("jetstream ended")]

     

       12
       +
           JetstreamEnded

     

       13
       +
       }

spacedust/src/lib.rs

···

       1
        
       pub mod consumer;

     

       0
        
       
     

       2
        
       pub mod server;

     

       3
        
       pub mod subscriber;

     

       4

···

       1
        
       pub mod consumer;

     

       2
       +
       pub mod error;

     

       3
        
       pub mod server;

     

       4
        
       pub mod subscriber;

     

       5

+23 -7

spacedust/src/main.rs

···

       3
        
       

     

       4
        
       use clap::Parser;

     

       5
        
       use tokio::sync::broadcast;

     

       0
        
       
     

       6
        
       

     

       7
        
       /// Aggregate links in the at-mosphere

     

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

     
···

       21
        
       

     

       22
        
       #[tokio::main]

     

       23
        
       async fn main() -> Result<(), String> {

     

       24
       -
           let args = Args::parse();

     

       25
        
       

     

       26
        
           // tokio broadcast keeps a single main output queue for all subscribers.

     

       27
        
           // each subscriber clones off a copy of an individual value for each recv.

     
···

       40
        
           // TODO: determine if a pathological case could blow this up (eg 1MB link

     

       41
        
           // paths + slow subscriber -> 16GiB queue)

     

       42
        
           let (b, _) = broadcast::channel(16_384);

     

       0
        
       
     

       43
        
       

     

       44
       -
           let consuming = consumer::consume(b.clone(), &args.jetstream, None, args.jetstream_no_zstd);

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       45
        
       

     

       46
       -
           let serving = server::serve(b);

     

       47
        
       

     

       48
       -
           tokio::select! {

     

       49
       -
               e = serving => eprintln!("serving failed: {e:?}"),

     

       50
       -
               e = consuming => eprintln!("consuming failed: {e:?}"),

     

       51
       -
           };

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       52
        
       

     

       53
        
           Ok(())

     

       54
        
       }

···

       3
        
       

     

       4
        
       use clap::Parser;

     

       5
        
       use tokio::sync::broadcast;

     

       6
       +
       use tokio_util::sync::CancellationToken;

     

       7
        
       

     

       8
        
       /// Aggregate links in the at-mosphere

     

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

     
···

       22
        
       

     

       23
        
       #[tokio::main]

     

       24
        
       async fn main() -> Result<(), String> {

     

       25
       +
           env_logger::init();

     

       26
        
       

     

       27
        
           // tokio broadcast keeps a single main output queue for all subscribers.

     

       28
        
           // each subscriber clones off a copy of an individual value for each recv.

     
···

       41
        
           // TODO: determine if a pathological case could blow this up (eg 1MB link

     

       42
        
           // paths + slow subscriber -> 16GiB queue)

     

       43
        
           let (b, _) = broadcast::channel(16_384);

     

       44
       +
           let consumer_sender = b.clone();

     

       45
        
       

     

       46
       +
           let shutdown = CancellationToken::new();

     

       47
       +
       

     

       48
       +
           let ctrlc_shutdown = shutdown.clone();

     

       49
       +
           ctrlc::set_handler(move || ctrlc_shutdown.cancel()).expect("failed to set ctrl-c handler");

     

       50
        
       

     

       51
       +
           let args = Args::parse();

     

       52
        
       

     

       53
       +
           let server_shutdown = shutdown.clone();

     

       54
       +
           let serving = tokio::spawn(async move {

     

       55
       +
               server::serve(b, server_shutdown).await

     

       56
       +
           });

     

       57
       +
       

     

       58
       +
           let consumer_shutdown = shutdown.clone();

     

       59
       +
           let consuming = tokio::spawn(async move {

     

       60
       +
               consumer::consume(consumer_sender, args.jetstream, None, args.jetstream_no_zstd, consumer_shutdown).await

     

       61
       +
           });

     

       62
       +
       

     

       63
       +
           let (served, consumed) = tokio::join!(serving, consuming);

     

       64
       +
           log::info!("serving ended: {served:?}");

     

       65
       +
           log::info!("consuming ended: {consumed:?}");

     

       66
       +
       

     

       67
       +
           log::info!("bye!");

     

       68
        
       

     

       69
        
           Ok(())

     

       70
        
       }

+23 -12

spacedust/src/server.rs

···

       1
       -
       use crate::subscriber;

     

       2
        
       use metrics::{histogram, counter};

     

       3
        
       use std::sync::Arc;

     

       4
        
       use crate::LinkEvent;

     
···

       19
        
       use tokio::sync::broadcast;

     

       20
        
       use tokio::time::Instant;

     

       21
        
       use tokio_tungstenite::tungstenite::protocol::Role;

     

       0
        
       
     

       22
        
       use async_trait::async_trait;

     

       23
        
       use std::collections::HashSet;

     

       24
        
       

     

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

     

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

     

       27
        
       

     

       28
       -
       pub async fn serve(b: broadcast::Sender<LinkEvent>) -> Result<(), String> {

     

       29
        
           let config_logging = ConfigLogging::StderrTerminal {

     

       30
        
               level: ConfigLoggingLevel::Info,

     

       31
        
           };

     
···

       58
        
               .map_err(|e| e.to_string())?,

     

       59
        
           );

     

       60
        
       

     

       61
       -
           let ctx = Context { spec, b };

     

       0
        
       
     

       62
        
       

     

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

     

       64
        
               .config(ConfigDropshot {

     
···

       68
        
               .start()

     

       69
        
               .map_err(|error| format!("failed to create server: {}", error))?;

     

       70
        
       

     

       71
       -
           server.await

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       72
        
       }

     

       73
        
       

     

       74
        
       #[derive(Debug, Clone)]

     

       75
        
       struct Context {

     

       76
        
           pub spec: Arc<serde_json::Value>,

     

       77
        
           pub b: broadcast::Sender<LinkEvent>,

     

       0
        
       
     

       78
        
       }

     

       79
        
       

     

       80
        
       async fn instrument_handler<T, H, R>(ctx: &RequestContext<T>, handler: H) -> Result<R, HttpError>

     
···

       266
        
           }

     

       267
        
       }

     

       268
        
       

     

       269
       -
       #[derive(Deserialize, JsonSchema)]

     

       270
       -
       struct QueryParams {

     

       271
       -
           _hello: Option<String>,

     

       272
       -
       }

     

       273
       -
       

     

       274
        
       #[channel {

     

       275
        
           protocol = WEBSOCKETS,

     

       276
        
           path = "/subscribe",

     

       277
        
       }]

     

       278
        
       async fn subscribe(

     

       279
       -
           ctx: RequestContext<Context>,

     

       280
        
           query: MultiSubscribeQuery,

     

       281
        
           upgraded: WebsocketConnection,

     

       282
        
       ) -> dropshot::WebsocketChannelResult {

     
···

       287
        
           )

     

       288
        
           .await;

     

       289
        
       

     

       290
       -
           let b = ctx.context().b.subscribe();

     

       0
        
       
     

       0
        
       
     

       291
        
       

     

       292
       -
           subscriber::subscribe(b, ws, query)

     

       0
        
       
     

       293
        
               .await

     

       294
        
               .map_err(|e| format!("boo: {e:?}"))?;

     

       295

···

       1
       +
       use crate::subscriber::Subscriber;

     

       2
        
       use metrics::{histogram, counter};

     

       3
        
       use std::sync::Arc;

     

       4
        
       use crate::LinkEvent;

     
···

       19
        
       use tokio::sync::broadcast;

     

       20
        
       use tokio::time::Instant;

     

       21
        
       use tokio_tungstenite::tungstenite::protocol::Role;

     

       22
       +
       use tokio_util::sync::CancellationToken;

     

       23
        
       use async_trait::async_trait;

     

       24
        
       use std::collections::HashSet;

     

       25
        
       

     

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

     

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

     

       28
        
       

     

       29
       +
       pub async fn serve(b: broadcast::Sender<LinkEvent>, shutdown: CancellationToken) -> Result<(), String> {

     

       30
        
           let config_logging = ConfigLogging::StderrTerminal {

     

       31
        
               level: ConfigLoggingLevel::Info,

     

       32
        
           };

     
···

       59
        
               .map_err(|e| e.to_string())?,

     

       60
        
           );

     

       61
        
       

     

       62
       +
           let sub_shutdown = shutdown.clone();

     

       63
       +
           let ctx = Context { spec, b, shutdown: sub_shutdown };

     

       64
        
       

     

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

     

       66
        
               .config(ConfigDropshot {

     
···

       70
        
               .start()

     

       71
        
               .map_err(|error| format!("failed to create server: {}", error))?;

     

       72
        
       

     

       73
       +
           tokio::select! {

     

       74
       +
               s = server.wait_for_shutdown() => {

     

       75
       +
                   log::error!("dropshot server ended: {s:?}");

     

       76
       +
                   s

     

       77
       +
               },

     

       78
       +
               _ = shutdown.cancelled() => {

     

       79
       +
                   log::info!("shutting down server");

     

       80
       +
                   server.close().await?;

     

       81
       +
                   Err("shutdown requested".to_string())

     

       82
       +
               }

     

       83
       +
           }

     

       84
        
       }

     

       85
        
       

     

       86
        
       #[derive(Debug, Clone)]

     

       87
        
       struct Context {

     

       88
        
           pub spec: Arc<serde_json::Value>,

     

       89
        
           pub b: broadcast::Sender<LinkEvent>,

     

       90
       +
           pub shutdown: CancellationToken,

     

       91
        
       }

     

       92
        
       

     

       93
        
       async fn instrument_handler<T, H, R>(ctx: &RequestContext<T>, handler: H) -> Result<R, HttpError>

     
···

       279
        
           }

     

       280
        
       }

     

       281
        
       

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       282
        
       #[channel {

     

       283
        
           protocol = WEBSOCKETS,

     

       284
        
           path = "/subscribe",

     

       285
        
       }]

     

       286
        
       async fn subscribe(

     

       287
       +
           reqctx: RequestContext<Context>,

     

       288
        
           query: MultiSubscribeQuery,

     

       289
        
           upgraded: WebsocketConnection,

     

       290
        
       ) -> dropshot::WebsocketChannelResult {

     
···

       295
        
           )

     

       296
        
           .await;

     

       297
        
       

     

       298
       +
           let Context { b, shutdown, .. } = reqctx.context();

     

       299
       +
           let sub_token = shutdown.child_token();

     

       300
       +
           let subscription = b.subscribe();

     

       301
        
       

     

       302
       +
           Subscriber::new(query, sub_token)

     

       303
       +
               .start(ws, subscription)

     

       304
        
               .await

     

       305
        
               .map_err(|e| format!("boo: {e:?}"))?;

     

       306

+142 -55

spacedust/src/subscriber.rs

···

       0
        
       
     

       0
        
       
     

       0
        
       
     

       1
        
       use crate::ClientEvent;

     

       2
        
       use crate::LinkEvent;

     

       3
        
       use crate::server::MultiSubscribeQuery;

     

       4
        
       use futures::SinkExt;

     

       5
        
       use std::error::Error;

     

       6
       -
       use tokio::sync::broadcast;

     

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

     

       0
        
       
     

       8
        
       use dropshot::WebsocketConnectionRaw;

     

       9
        
       

     

       10
       -
       pub async fn subscribe(

     

       11
       -
           mut sub: broadcast::Receiver<LinkEvent>,

     

       12
       -
           mut ws: WebSocketStream<WebsocketConnectionRaw>,

     

       13
        
           query: MultiSubscribeQuery,

     

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

     

       15
       -
           // TODO: pingpong

     

       16
        
       

     

       17
       -
           loop {

     

       18
       -
               match sub.recv().await {

     

       19
       -
                   Ok(link) => {

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       20
        
       

     

       21
       -
                       // subject + subject DIDs are logical OR

     

       22
       -
                       let target_did = if link.target.starts_with("did:") {

     

       23
       -
                           link.target.clone()

     

       24
       -
                       } else {

     

       25
       -
                           let Some(rest) = link.target.strip_prefix("at://") else {

     

       26
       -
                               continue;

     

       27
       -
                           };

     

       28
       -
                           if let Some((did, _)) = rest.split_once("/") {

     

       29
       -
                               did

     

       30
       -
                           } else {

     

       31
       -
                               rest

     

       32
       -
                           }.to_string()

     

       33
       -
                       };

     

       34
       -
                       if !(query.wanted_subjects.contains(&link.target) || query.wanted_subject_dids.contains(&target_did) || query.wanted_subjects.is_empty() && query.wanted_subject_dids.is_empty()) {

     

       35
       -
                           // wowwww ^^ fix that

     

       36
       -
                           continue;

     

       37
       -
                       }

     

       38
        
       

     

       39
       -
                       // subjects together with sources are logical AND

     

       40
        
       

     

       41
       -
                       if !query.wanted_sources.is_empty() {

     

       42
       -
                           let undotted = link.path.strip_prefix('.').unwrap_or_else(|| {

     

       43
       -
                               eprintln!("link path did not have expected '.' prefix: {}", link.path);

     

       44
       -
                               ""

     

       45
       -
                           });

     

       46
       -
                           let source = format!("{}:{undotted}", link.collection);

     

       47
       -
                           if !query.wanted_sources.contains(&source) {

     

       48
       -
                               continue;

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       49
        
                           }

     

       50
        
                       }

     

       51
       -
       

     

       52
       -
                       let ev = ClientEvent {

     

       53
       -
                           kind: "link".to_string(),

     

       54
       -
                           origin: "live".to_string(),

     

       55
       -
                           link: link.into(),

     

       56
       -
                       };

     

       57
       -
                       let json = serde_json::to_string(&ev)?;

     

       58
       -
                       if let Err(e) = ws.send(Message::Text(json.into())).await {

     

       59
       -
                           eprintln!("client: failed to send event: {e:?}");

     

       60
       -
                           ws.close(None).await?; // TODO: do we need this one??

     

       61
        
                           break;

     

       62
       -
                       }

     

       63
       -
                   }

     

       64
       -
                   Err(broadcast::error::RecvError::Closed) => {

     

       65
       -
                       ws.close(None).await?; // TODO: send reason

     

       66
       -
                       break;

     

       67
        
                   }

     

       68
       -
                   Err(broadcast::error::RecvError::Lagged(_n_missed)) => {

     

       69
       -
                       eprintln!("client lagged, closing");

     

       70
       -
                       ws.close(None).await?; // TODO: send reason

     

       71
       -
                       break;

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       72
        
                   }

     

       73
        
               }

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       74
        
           }

     

       75
       -
           Ok(())

     

       76
        
       }

···

       1
       +
       use tokio::time::interval;

     

       2
       +
       use std::time::Duration;

     

       3
       +
       use futures::StreamExt;

     

       4
        
       use crate::ClientEvent;

     

       5
        
       use crate::LinkEvent;

     

       6
        
       use crate::server::MultiSubscribeQuery;

     

       7
        
       use futures::SinkExt;

     

       8
        
       use std::error::Error;

     

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

     

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

     

       11
       +
       use tokio_util::sync::CancellationToken;

     

       12
        
       use dropshot::WebsocketConnectionRaw;

     

       13
        
       

     

       14
       +
       const PING_PERIOD: Duration = Duration::from_secs(30);

     

       15
       +
       

     

       16
       +
       pub struct Subscriber {

     

       17
        
           query: MultiSubscribeQuery,

     

       18
       +
           shutdown: CancellationToken,

     

       19
       +
       }

     

       20
        
       

     

       21
       +
       impl Subscriber {

     

       22
       +
           pub fn new(

     

       23
       +
               query: MultiSubscribeQuery,

     

       24
       +
               shutdown: CancellationToken,

     

       25
       +
           ) -> Self {

     

       26
       +
               log::warn!("new sub...");

     

       27
       +
               Self { query, shutdown }

     

       28
       +
           }

     

       29
        
       

     

       30
       +
           pub async fn start(

     

       31
       +
               self,

     

       32
       +
               ws: WebSocketStream<WebsocketConnectionRaw>,

     

       33
       +
               mut receiver: broadcast::Receiver<LinkEvent>

     

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

     

       35
       +
               log::warn!("starting new sub...");

     

       36
       +
               let mut ping_state = None;

     

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

     

       38
       +
               let mut ping_interval = interval(PING_PERIOD);

     

       39
       +
               let _guard = self.shutdown.clone().drop_guard();

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       40
        
       

     

       41
       +
               // TODO: do we need to timeout ws sends??

     

       42
        
       

     

       43
       +
               loop {

     

       44
       +
                   tokio::select! {

     

       45
       +
                       l = receiver.recv() => match l {

     

       46
       +
                           Ok(link) => if let Some(message) = self.filter(link) {

     

       47
       +
                               if let Err(e) = ws_sender.send(message).await {

     

       48
       +
                                   log::warn!("failed to send link, dropping subscriber: {e:?}");

     

       49
       +
                                   break;

     

       50
       +
                               }

     

       51
       +
                           },

     

       52
       +
                           Err(RecvError::Closed) => self.shutdown.cancel(),

     

       53
       +
                           Err(RecvError::Lagged(n)) => {

     

       54
       +
                               log::warn!("dropping lagging subscriber (missed {n} messages already)");

     

       55
       +
                               self.shutdown.cancel();

     

       56
       +
                           }

     

       57
       +
                       },

     

       58
       +
                       cm = ws_receiver.next() => match cm {

     

       59
       +
                           Some(Ok(Message::Ping(state))) => {

     

       60
       +
                               if let Err(e) = ws_sender.send(Message::Pong(state)).await {

     

       61
       +
                                   log::error!("failed to reply pong to subscriber: {e:?}");

     

       62
       +
                                   break;

     

       63
       +
                               }

     

       64
       +
                           }

     

       65
       +
                           Some(Ok(Message::Pong(state))) => {

     

       66
       +
                               if let Some(expected_state) = ping_state {

     

       67
       +
                                   if *state == expected_state {

     

       68
       +
                                       ping_state = None; // good

     

       69
       +
                                   } else {

     

       70
       +
                                       log::error!("subscriber returned a pong with the wrong state, dropping");

     

       71
       +
                                       self.shutdown.cancel();

     

       72
       +
                                   }

     

       73
       +
                               } else {

     

       74
       +
                                   log::error!("subscriber sent a pong when none was expected");

     

       75
       +
                                   self.shutdown.cancel();

     

       76
       +
                               }

     

       77
       +
                           }

     

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

     

       79
       +
                           Some(Err(e)) => {

     

       80
       +
                               log::error!("failed to receive subscriber message: {e:?}");

     

       81
       +
                               break;

     

       82
       +
                           }

     

       83
       +
                           None => {

     

       84
       +
                               log::trace!("end of subscriber messages. bye!");

     

       85
       +
                               break;

     

       86
       +
                           }

     

       87
       +
                       },

     

       88
       +
                       _ = ping_interval.tick() => {

     

       89
       +
                           if ping_state.is_some() {

     

       90
       +
                               log::warn!("did not recieve pong within {PING_PERIOD:?}, dropping subscriber");

     

       91
       +
                               self.shutdown.cancel();

     

       92
       +
                           } else {

     

       93
       +
                               let new_state: [u8; 8] = rand::random();

     

       94
       +
                               let ping = new_state.to_vec().into();

     

       95
       +
                               ping_state = Some(new_state);

     

       96
       +
                               if let Err(e) = ws_sender.send(Message::Ping(ping)).await {

     

       97
       +
                                   log::error!("failed to send ping to subscriber, dropping: {e:?}");

     

       98
       +
                                   self.shutdown.cancel();

     

       99
       +
                               }

     

       100
        
                           }

     

       101
        
                       }

     

       102
       +
                       _ = self.shutdown.cancelled() => {

     

       103
       +
                           log::info!("subscriber shutdown requested, bye!");

     

       104
       +
                           if let Err(e) = ws_sender.close().await {

     

       105
       +
                               log::warn!("failed to close subscriber: {e:?}");

     

       106
       +
                           }

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       107
        
                           break;

     

       108
       +
                       },

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       109
        
                   }

     

       110
       +
               }

     

       111
       +
               log::trace!("end of subscriber. bye!");

     

       112
       +
               Ok(())

     

       113
       +
           }

     

       114
       +
       

     

       115
       +
           fn filter(

     

       116
       +
               &self,

     

       117
       +
               link: LinkEvent,

     

       118
       +
               // mut sender: impl Sink<Message> + Unpin

     

       119
       +
           ) -> Option<Message> {

     

       120
       +
               let query = &self.query;

     

       121
       +
       

     

       122
       +
               // subject + subject DIDs are logical OR

     

       123
       +
               let target_did = if link.target.starts_with("did:") {

     

       124
       +
                   link.target.clone()

     

       125
       +
               } else {

     

       126
       +
                   let Some(rest) = link.target.strip_prefix("at://") else {

     

       127
       +
                       return None

     

       128
       +
                   };

     

       129
       +
                   if let Some((did, _)) = rest.split_once("/") {

     

       130
       +
                       did

     

       131
       +
                   } else {

     

       132
       +
                       rest

     

       133
       +
                   }.to_string()

     

       134
       +
               };

     

       135
       +
               if !(query.wanted_subjects.contains(&link.target) || query.wanted_subject_dids.contains(&target_did) || query.wanted_subjects.is_empty() && query.wanted_subject_dids.is_empty()) {

     

       136
       +
                   // wowwww ^^ fix that

     

       137
       +
                   return None

     

       138
       +
               }

     

       139
       +
       

     

       140
       +
               // subjects together with sources are logical AND

     

       141
       +
       

     

       142
       +
               if !query.wanted_sources.is_empty() {

     

       143
       +
                   let undotted = link.path.strip_prefix('.').unwrap_or_else(|| {

     

       144
       +
                       eprintln!("link path did not have expected '.' prefix: {}", link.path);

     

       145
       +
                       ""

     

       146
       +
                   });

     

       147
       +
                   let source = format!("{}:{undotted}", link.collection);

     

       148
       +
                   if !query.wanted_sources.contains(&source) {

     

       149
       +
                       return None

     

       150
        
                   }

     

       151
        
               }

     

       152
       +
       

     

       153
       +
               let ev = ClientEvent {

     

       154
       +
                   kind: "link".to_string(),

     

       155
       +
                   origin: "live".to_string(),

     

       156
       +
                   link: link.into(),

     

       157
       +
               };

     

       158
       +
       

     

       159
       +
               let json = serde_json::to_string(&ev).unwrap();

     

       160
       +
       

     

       161
       +
               Some(Message::Text(json.into()))

     

       162
        
           }

     

       0
        
       
     

       163
        
       }