commit 85f997c43dc9d587130dc4a175b0af14ab9abbb9 · ptr.pet/nsid-tracker

.gitignore

···

       28
       28
        
       events.sqlite-shm

     

       29
       29
        
       

     

       30
       30
        
       result

     

       31
       31
       +
       server/bsky_zstd_dictionary

server/src/bsky_zstd_dictionary

This is a binary file and will not be displayed.

+40 -5

server/src/db.rs

···

       1
       1
       +
       use atproto_jetstream::JetstreamEvent;

     

       1
       2
        
       use fjall::{Config, Keyspace, Partition, PartitionCreateOptions};

     

       2
       3
        
       use rkyv::{Archive, Deserialize, Serialize, rancor::Error};

     

       3
       4
        
       use smol_str::SmolStr;

     
···

       19
       20
        
           pub deleted: bool,

     

       20
       21
        
       }

     

       21
       22
        
       

     

       23
       23
       +
       pub struct EventRecord {

     

       24
       24
       +
           nsid: SmolStr,

     

       25
       25
       +
           timestamp: u64,

     

       26
       26
       +
           deleted: bool,

     

       27
       27
       +
       }

     

       28
       28
       +
       

     

       29
       29
       +
       impl EventRecord {

     

       30
       30
       +
           pub fn from_jetstream(event: JetstreamEvent) -> Option<Self> {

     

       31
       31
       +
               match event {

     

       32
       32
       +
                   JetstreamEvent::Commit {

     

       33
       33
       +
                       time_us, commit, ..

     

       34
       34
       +
                   } => Some(Self {

     

       35
       35
       +
                       nsid: commit.collection.into(),

     

       36
       36
       +
                       timestamp: time_us,

     

       37
       37
       +
                       deleted: false,

     

       38
       38
       +
                   }),

     

       39
       39
       +
                   JetstreamEvent::Delete {

     

       40
       40
       +
                       time_us, commit, ..

     

       41
       41
       +
                   } => Some(Self {

     

       42
       42
       +
                       nsid: commit.collection.into(),

     

       43
       43
       +
                       timestamp: time_us,

     

       44
       44
       +
                       deleted: true,

     

       45
       45
       +
                   }),

     

       46
       46
       +
                   _ => None,

     

       47
       47
       +
               }

     

       48
       48
       +
           }

     

       49
       49
       +
       }

     

       50
       50
       +
       

     

       22
       51
        
       // counts is nsid -> NsidCounts

     

       23
       52
        
       // hits is tree per nsid: timestamp -> NsidHit

     

       24
       53
        
       pub struct Db {

     
···

       57
       86
        
               }))

     

       58
       87
        
           }

     

       59
       88
        
       

     

       60
       60
       -
           pub fn record_event(&self, nsid: &str, timestamp: u64, deleted: bool) -> AppResult<()> {

     

       61
       61
       -
               self.insert_event(nsid, timestamp, deleted)?;

     

       89
       89
       +
           pub fn record_event(&self, e: EventRecord) -> AppResult<()> {

     

       90
       90
       +
               let EventRecord {

     

       91
       91
       +
                   nsid,

     

       92
       92
       +
                   timestamp,

     

       93
       93
       +
                   deleted,

     

       94
       94
       +
               } = e;

     

       95
       95
       +
       

     

       96
       96
       +
               self.insert_event(&nsid, timestamp, deleted)?;

     

       62
       97
        
               // increment count

     

       63
       63
       -
               let mut counts = self.get_count(nsid)?;

     

       98
       98
       +
               let mut counts = self.get_count(&nsid)?;

     

       64
       99
        
               counts.last_seen = timestamp;

     

       65
       100
        
               if deleted {

     

       66
       101
        
                   counts.deleted_count += 1;

     

       67
       102
        
               } else {

     

       68
       103
        
                   counts.count += 1;

     

       69
       104
        
               }

     

       70
       70
       -
               self.insert_count(nsid, counts.clone())?;

     

       105
       105
       +
               self.insert_count(&nsid, counts.clone())?;

     

       71
       106
        
               if self.event_broadcaster.receiver_count() > 0 {

     

       72
       72
       -
                   let _ = self.event_broadcaster.send((SmolStr::new(nsid), counts));

     

       107
       107
       +
                   let _ = self.event_broadcaster.send((SmolStr::new(&nsid), counts));

     

       73
       108
        
               }

     

       74
       109
        
               Ok(())

     

       75
       110
        
           }

+46 -29

server/src/main.rs

···

       1
       1
        
       use std::sync::Arc;

     

       2
       2
        
       

     

       3
       3
        
       use atproto_jetstream::{CancellationToken, Consumer, EventHandler, JetstreamEvent};

     

       4
       4
       +
       use tokio::sync::mpsc::{Receiver, Sender};

     

       4
       5
        
       

     

       5
       5
       -
       use crate::{api::serve, db::Db};

     

       6
       6
       +
       use crate::{

     

       7
       7
       +
           api::serve,

     

       8
       8
       +
           db::{Db, EventRecord},

     

       9
       9
       +
       };

     

       6
       10
        
       

     

       7
       11
        
       mod api;

     

       8
       12
        
       mod db;

     

       9
       13
        
       mod error;

     

       14
       14
       +
       

     

       15
       15
       +
       const BSKY_ZSTD_DICT: &[u8] = include_bytes!("./bsky_zstd_dictionary");

     

       10
       16
        
       

     

       11
       17
        
       struct JetstreamHandler {

     

       12
       12
       -
           db: Arc<Db>,

     

       18
       18
       +
           tx: Sender<EventRecord>,

     

       19
       19
       +
       }

     

       20
       20
       +
       

     

       21
       21
       +
       impl JetstreamHandler {

     

       22
       22
       +
           fn new() -> (Self, Receiver<EventRecord>) {

     

       23
       23
       +
               let (tx, rx) = tokio::sync::mpsc::channel(1000);

     

       24
       24
       +
               (Self { tx }, rx)

     

       25
       25
       +
           }

     

       13
       26
        
       }

     

       14
       27
        
       

     

       15
       28
        
       #[async_trait::async_trait]

     

       16
       29
        
       impl EventHandler for JetstreamHandler {

     

       17
       30
        
           async fn handle_event(&self, event: JetstreamEvent) -> anyhow::Result<()> {

     

       18
       18
       -
               let db = self.db.clone();

     

       19
       19
       -
               tokio::task::spawn_blocking(move || {

     

       20
       20
       -
                   let result = match event {

     

       21
       21
       -
                       JetstreamEvent::Commit {

     

       22
       22
       -
                           time_us, commit, ..

     

       23
       23
       -
                       } => db.record_event(&commit.collection, time_us, false),

     

       24
       24
       -
                       JetstreamEvent::Delete {

     

       25
       25
       -
                           time_us, commit, ..

     

       26
       26
       -
                       } => db.record_event(&commit.collection, time_us, true),

     

       27
       27
       -
                       _ => Ok(()),

     

       28
       28
       -
                   };

     

       29
       29
       -
                   if let Err(err) = result {

     

       30
       30
       -
                       tracing::error!("couldn't record event: {err}");

     

       31
       31
       -
                   }

     

       32
       32
       -
               });

     

       31
       31
       +
               if let Some(e) = EventRecord::from_jetstream(event) {

     

       32
       32
       +
                   self.tx.send(e).await?;

     

       33
       33
       +
               }

     

       33
       34
        
               Ok(())

     

       34
       35
        
           }

     

       35
       36
        
       

     
···

       40
       41
        
       

     

       41
       42
        
       #[tokio::main]

     

       42
       43
        
       async fn main() {

     

       43
       43
       -
           tracing_subscriber::fmt::init();

     

       44
       44
       +
           tracing_subscriber::fmt::fmt().compact().init();

     

       44
       45
        
       

     

       45
       46
        
           let db = Arc::new(Db::new().expect("couldnt create db"));

     

       46
       47
        
       

     

       47
       47
       -
           let consumer = Consumer::new(atproto_jetstream::ConsumerTaskConfig {

     

       48
       48
       -
               compression: false,

     

       48
       48
       +
           tokio::fs::write("./bsky_zstd_dictionary", BSKY_ZSTD_DICT)

     

       49
       49
       +
               .await

     

       50
       50
       +
               .expect("could not write bsky zstd dict");

     

       51
       51
       +
       

     

       52
       52
       +
           let jetstream = Consumer::new(atproto_jetstream::ConsumerTaskConfig {

     

       53
       53
       +
               compression: true,

     

       49
       54
        
               jetstream_hostname: "jetstream2.us-west.bsky.network".into(),

     

       50
       55
        
               collections: Vec::new(),

     

       51
       56
        
               dids: Vec::new(),

     

       52
       57
        
               max_message_size_bytes: None,

     

       53
       58
        
               cursor: None,

     

       54
       59
        
               require_hello: true,

     

       55
       55
       -
               zstd_dictionary_location: String::new(),

     

       60
       60
       +
               zstd_dictionary_location: "./bsky_zstd_dictionary".into(),

     

       56
       61
        
               user_agent: "nsid-tracker/0.0.1".into(),

     

       57
       62
        
           });

     

       58
       63
        
       

     

       59
       59
       -
           tracing::info!("running jetstream consumer...");

     

       64
       64
       +
           let (event_handler, mut event_rx) = JetstreamHandler::new();

     

       65
       65
       +
       

     

       60
       66
        
           let cancel_token = CancellationToken::new();

     

       61
       61
       -
           tokio::spawn({

     

       67
       67
       +
           tokio::spawn(async move {

     

       68
       68
       +
               jetstream

     

       69
       69
       +
                   .register_handler(Arc::new(event_handler))

     

       70
       70
       +
                   .await

     

       71
       71
       +
                   .expect("cant register handler");

     

       72
       72
       +
               jetstream

     

       73
       73
       +
                   .run_background(cancel_token.clone())

     

       74
       74
       +
                   .await

     

       75
       75
       +
                   .expect("cant run jetstream");

     

       76
       76
       +
           });

     

       77
       77
       +
       

     

       78
       78
       +
           std::thread::spawn({

     

       62
       79
        
               let db = db.clone();

     

       63
       63
       -
               async move {

     

       64
       64
       -
                   consumer

     

       65
       65
       -
                       .register_handler(Arc::new(JetstreamHandler { db }))

     

       66
       66
       -
                       .await

     

       67
       67
       -
                       .unwrap();

     

       68
       68
       -
                   consumer.run_background(cancel_token.clone()).await.unwrap();

     

       80
       80
       +
               move || {

     

       81
       81
       +
                   while let Some(e) = event_rx.blocking_recv() {

     

       82
       82
       +
                       if let Err(e) = db.record_event(e) {

     

       83
       83
       +
                           tracing::error!("failed to record event: {}", e);

     

       84
       84
       +
                       }

     

       85
       85
       +
                   }

     

       69
       86
        
               }

     

       70
       87
        
           });

     

       71
       88