relay-legacy/src/lib.rs

83 lines
2.5 KiB
Rust

use std::collections::BTreeSet;
use anyhow::Result;
use bytes::Bytes;
use tokio::sync::{broadcast, mpsc, Mutex};
use wire_proto::StreamEvent;
pub struct RelayServer {
pub db: sled::Db,
pub db_history: sled::Tree,
pub db_users: sled::Tree,
pub db_index_cursors: sled::Tree,
pub known_good_hosts: Mutex<BTreeSet<String>>,
pub active_indexers: Mutex<BTreeSet<String>>,
pub event_tx: mpsc::Sender<StreamEvent>,
pub raw_block_tx: broadcast::Sender<Bytes>,
}
impl RelayServer {
pub fn new(db: sled::Db, event_tx: mpsc::Sender<StreamEvent>) -> Self {
let (raw_block_tx, _) = broadcast::channel(128);
let hosts = db
.get("hosts")
.expect("Failed to read db entry for hosts")
.and_then(|v| serde_ipld_dagcbor::from_slice::<Vec<String>>(&v).ok())
.unwrap_or_default();
Self {
event_tx,
raw_block_tx,
known_good_hosts: Mutex::new(hosts.into_iter().collect()),
active_indexers: Default::default(),
db_history: db
.open_tree("history")
.expect("failed to open history tree"),
db_users: db.open_tree("users").expect("failed to open users tree"),
db_index_cursors: db
.open_tree("index_cursors")
.expect("failed to ope index_cursors tree"),
db,
}
}
pub async fn add_good_host(&self, host: String) -> Result<()> {
let mut hosts = self.known_good_hosts.lock().await;
if hosts.contains(&host) {
return Ok(());
}
tracing::debug!(%host, "discovered new known-good host");
hosts.insert(host);
let hosts_copy = hosts.iter().collect::<Vec<_>>();
let serialized_hosts = serde_ipld_dagcbor::to_vec(&hosts_copy)?;
drop(hosts);
self.db.insert("hosts", serialized_hosts)?;
Ok(())
}
pub async fn remove_good_host(&self, host: String) -> Result<()> {
tracing::debug!(%host, "dropping known-good host");
let mut hosts = self.known_good_hosts.lock().await;
hosts.remove(&host);
let hosts_copy = hosts.iter().collect::<Vec<_>>();
let serialized_hosts = serde_ipld_dagcbor::to_vec(&hosts_copy)?;
drop(hosts);
self.db.insert("hosts", serialized_hosts)?;
Ok(())
}
}
pub mod http;
pub mod indexer;
pub mod relay_subscription;
pub mod request_crawl;
pub mod sequencer;
pub mod tls;
pub mod user;
pub mod wire_proto;