Repository: libsql/libsql
File: libsql-server/src/broadcaster.rs (Rust)

Code before:
use std::collections::HashMap;

use bytes::Bytes;
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use tokio::sync::broadcast;

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct TableName(Bytes);

#[derive(Debug, Default)]
pub struct UpdateBroadcaster {
    senders: Mutex<HashMap<String, broadcast::Sender<UpdateHookMsg>>>,
}

#[derive(Debug, Clone, Serialize)]
pub struct UpdateHookMsg {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub action: Option<Action>,
    pub rowid: i64,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum Action {
    UNKNOWN,
    DELETE,
    INSERT,
    UPDATE,
}

impl From<rusqlite::hooks::Action> for Action {
    fn from(value: rusqlite::hooks::Action) -> Self {
        match value {
            rusqlite::hooks::Action::SQLITE_DELETE => Action::DELETE,
            rusqlite::hooks::Action::SQLITE_INSERT => Action::INSERT,
            rusqlite::hooks::Action::SQLITE_UPDATE => Action::UPDATE,
            _ => Action::UNKNOWN,
        }
    }
}

impl From<&str> for Action {
    fn from(value: &str) -> Self {
        match value {
            "delete" => Action::DELETE,
            "insert" => Action::INSERT,
            "update" => Action::UPDATE,
            _ => Action::UNKNOWN,
        }
    }
}

impl UpdateBroadcaster {
    pub fn new() -> Self {
        UpdateBroadcaster {
            senders: Mutex::new(HashMap::new()),
        }
    }

    pub fn subscribe(&self, table: &str) -> broadcast::Receiver<UpdateHookMsg> {
        let mut senders = self.senders.lock();
        if let Some(sender) = senders.get(table) {
            return sender.subscribe();
        }

        let (sender, receiver) = broadcast::channel(16);
        senders.insert(table.to_owned(), sender);
        receiver
    }

    pub fn notify(&self, action: rusqlite::hooks::Action, table: &str, rowid: i64) {
        let mut senders = self.senders.lock();
        let readers_gone = senders.get_mut(table).map_or(false, |sender| {
            let msg = UpdateHookMsg {
                action: Some(action.into()),
                rowid,
            };
            match sender.send(msg) {
                Err(_) => return true,
                _ => return false,
            }
        });

        if readers_gone {
            senders.remove(table);
        }
    }
}


Review comment (suggestion):
I think this should be handled by the subscriber. Either with an `unsubscribe` method, or better yet, have `subscribe` return a `Subscription` object that unsubscribed on drop. If the [subscribed count](https://docs.rs/tokio/latest/tokio/sync/broadcast/struct.Sender.html#method.receiver_count) drop to 0, remove the channel from the map.

Code after:
use std::{
    collections::{hash_map::Entry, BTreeMap, HashMap},
    mem,
    sync::Arc,
};

use parking_lot::Mutex;
use serde::Serialize;
use tokio::sync::broadcast::{self};
use tokio_stream::wrappers::BroadcastStream;

#[derive(Debug, Copy, Clone, Serialize, Default)]
pub struct BroadcastMsg {
    #[serde(skip_serializing_if = "is_zero")]
    pub unknown: u64,
    #[serde(skip_serializing_if = "is_zero")]
    pub delete: u64,
    #[serde(skip_serializing_if = "is_zero")]
    pub insert: u64,
    #[serde(skip_serializing_if = "is_zero")]
    pub update: u64,
}

#[inline(always)]
fn is_zero(num: &u64) -> bool {
    *num == 0
}

#[derive(Debug, Default)]
pub struct BroadcasterInner {
    state: Mutex<BTreeMap<String, BroadcastMsg>>,
    senders: Mutex<HashMap<String, broadcast::Sender<BroadcastMsg>>>,
}

#[derive(Debug, Default, Clone)]
pub struct Broadcaster {
    inner: Arc<BroadcasterInner>,
}

impl Broadcaster {
    const BROADCAST_CAP: usize = 1024;

    pub fn notify(&self, table: &str, action: rusqlite::hooks::Action) {
        let mut state = self.inner.state.lock();
        if let Some(entry) = state.get_mut(table) {
            Self::increment(entry, action);
        } else {
            let mut entry = BroadcastMsg::default();
            Self::increment(&mut entry, action);
            state.insert(table.into(), entry);
        }
    }

    fn increment(value: &mut BroadcastMsg, action: rusqlite::hooks::Action) {
        match action {
            rusqlite::hooks::Action::SQLITE_DELETE => value.delete += 1,
            rusqlite::hooks::Action::SQLITE_INSERT => value.insert += 1,
            rusqlite::hooks::Action::SQLITE_UPDATE => value.update += 1,
            _ => value.unknown += 1,
        }
    }

    pub fn commit(&self) {
        let senders = self.inner.senders.lock();
        let mut state = self.inner.state.lock();
        let state = mem::take(&mut *state);

        state.into_iter().for_each(|(table, entry)| {
            if let Some(sender) = senders.get(&table) {
                let _ = sender.send(entry);
            }
        });
    }

    pub fn rollback(&self) {
        self.inner.state.lock().clear();
    }

    pub fn subscribe(&self, table: String) -> BroadcastStream<BroadcastMsg> {
        let receiver = match self.inner.senders.lock().entry(table) {
            Entry::Occupied(entry) => entry.get().subscribe(),
            Entry::Vacant(entry) => {
                let (sender, receiver) = broadcast::channel(Self::BROADCAST_CAP);
                entry.insert(sender);
                receiver
            }
        };

        BroadcastStream::new(receiver)
    }

    pub fn unsubscribe(&self, table: &String) -> bool {
        let mut tables = self.inner.senders.lock();
        if let Some(sender) = tables.get(table) {
            if sender.receiver_count() == 0 {
                tables.remove(table);
                if tables.is_empty() {
                    return false;
                }
            }
        }
        return true;
    }
}