ronantakizawa/codereview-bench huggingface dataset:
Repository: libsql/libsql
File: libsql-server/src/http/user/listen.rs (Rust)

Code before:
    query: Query<ListenQuery>,
) -> crate::Result<axum::body::StreamBody<ReaderStream<DuplexStream>>> {
    let namespace = namespace_from_headers(
        &headers,
        state.disable_default_namespace,
        state.disable_namespaces,
    )?;

    if !auth.is_namespace_authorized(&namespace) {
        return Err(Error::NamespaceDoesntExist(namespace.to_string()));
    }

    // TODO: validate table
    let update_broadcaster = state.namespaces.update_broadcaster(namespace).await?;
    let mut rx = update_broadcaster.subscribe(query.table.as_ref());
    let action = query.action.clone();

    let (reader, writer) = tokio::io::duplex(8 * 1024);

    BLOCKING_RT.spawn_blocking(|| {
        futures::executor::block_on(async move {
            let mut writer = SyncIoBridge::new(writer);
            // TODO: stop when request is cancelled
            while recv_listen_msg(&mut writer, &mut rx, &action).await {}
        })
    });

    let stream = tokio_util::io::ReaderStream::new(reader);
    let stream = axum::body::StreamBody::new(stream);
    Ok(stream)
}

async fn recv_listen_msg(
    writer: &mut SyncIoBridge<DuplexStream>,
    rx: &mut tokio::sync::broadcast::Receiver<UpdateHookMsg>,
    action: &Option<Action>,
) -> bool {
    let keep_listening = match rx.recv().await {
        Ok(msg) => {
            send_msg_if_matches(writer, msg, action);
            true
        }
        Err(error) => {
            handle_listen_error(writer, ListenError::Error(error.clone()));
            error != RecvError::Closed
        }
    };
    keep_listening
}

fn handle_listen_error(writer: &mut SyncIoBridge<DuplexStream>, error: ListenError) {

Review comment (suggestion):
I think all of that would be better handled by: https://docs.rs/axum-extra/0.9.3/axum_extra/json_lines/struct.JsonLines.html

Code after:
#[derive(Deserialize)]
pub struct ListenQuery {
    table: String,
    action: Vec<Action>,
}

pub(super) async fn handle_listen(
    auth: Authenticated,
    AxumState(state): AxumState<AppState>,
    headers: HeaderMap,
    uri: Uri,
    query: Query<ListenQuery>,
) -> crate::Result<Response> {
    let namespace = namespace_from_headers(
        &headers,
        state.disable_default_namespace,
        state.disable_namespaces,
    )?;

    if !auth.is_namespace_authorized(&namespace) {
        return Err(Error::NamespaceDoesntExist(namespace.to_string()));
    }

    if let Some(primary_url) = state.primary_url {
        return Ok(Response::builder()
            .status(307)
            .header("Location", primary_url + uri.path())
            .body(BoxBody::default())
            .unwrap());
    }

    let stream = listen_stream(
        state.namespaces.clone(),
        namespace,
        query.table.clone(),
        query.action.clone(),
    )
    .await;
    Ok(JsonLines::new(stream).into_response())
}

static LAGGED_MSG: &str = "some changes were lost";

#[derive(Debug, Serialize)]
#[serde(rename_all = "snake_case")]
enum AggregatorEvent {
    Error(&'static str),
    #[serde(untagged)]
    Changes(BroadcastMsg),
}
Repository: libsql/libsql
File: libsql-server/src/namespace/store.rs (Rust)

Code before:
            }
        }

        while let Some(_) = set.join_next().await.transpose()?.transpose()? {}

        self.inner.metadata.shutdown().await?;
        self.inner.store.invalidate_all();
        self.inner.store.run_pending_tasks().await;
        Ok(())
    }

    pub(crate) async fn stats(&self, namespace: NamespaceName) -> crate::Result<Arc<Stats>> {
        self.with(namespace, |ns| ns.stats.clone()).await
    }

    pub(crate) async fn subscribe(
        &self,
        namespace: NamespaceName,
        table: String,
    ) -> crate::Result<UpdateSubscription> {
        self.with(namespace, |ns| ns.broadcaster.subscribe(table))
            .await
    }

    pub(crate) async fn unsubscribe(&self, namespace: NamespaceName, table: String) {
        if self.inner.store.contains_key(&namespace) {
            _ = self
                .with(namespace, |ns| ns.broadcaster.unsubscribe(table))
                .await;
        }
    }

    pub(crate) async fn config_store(
        &self,
        namespace: NamespaceName,
    ) -> crate::Result<MetaStoreHandle> {
        self.with(namespace, |ns| ns.db_config_store.clone()).await
    }

    pub(crate) fn meta_store(&self) -> &MetaStore {
        &self.inner.metadata
    }

    pub(crate) fn schema_locks(&self) -> &SchemaLocksRegistry {
        &self.inner.schema_locks
    }
}


Review comment (suggestion):
@LucioFranco @MarinPostma, I'm not happy with the workaround. I used to unsubscribe only if the namespace was still loaded. We should probably have a function like `with` that only runs if the namespace is loaded

Code after:
                set.spawn(async move {
                    ns.shutdown(snapshow_at_shutdown).await?;
                    Ok::<_, anyhow::Error>(())
                });
            }
        }

        while let Some(_) = set.join_next().await.transpose()?.transpose()? {}

        self.inner.metadata.shutdown().await?;
        self.inner.store.invalidate_all();
        self.inner.store.run_pending_tasks().await;
        Ok(())
    }

    pub(crate) async fn stats(&self, namespace: NamespaceName) -> crate::Result<Arc<Stats>> {
        self.with(namespace, |ns| ns.stats.clone()).await
    }

    pub(crate) fn broadcaster(&self, namespace: NamespaceName) -> BroadcasterHandle {
        self.inner.broadcasters.handle(namespace)
    }

    pub(crate) fn subscribe(
        &self,
        namespace: NamespaceName,
        table: String,
    ) -> BroadcastStream<BroadcastMsg> {
        self.inner.broadcasters.subscribe(namespace, table)
    }

    pub(crate) fn unsubscribe(&self, namespace: NamespaceName, table: &String) {
        self.inner.broadcasters.unsubscribe(namespace, table);
    }

    pub(crate) async fn config_store(
        &self,
        namespace: NamespaceName,
    ) -> crate::Result<MetaStoreHandle> {
        self.with(namespace, |ns| ns.db_config_store.clone()).await
    }

    pub(crate) fn meta_store(&self) -> &MetaStore {
        &self.inner.metadata
    }

    pub(crate) fn schema_locks(&self) -> &SchemaLocksRegistry {
        &self.inner.schema_locks
    }
}
Repository: libsql/libsql
File: libsql-server/src/broadcaster.rs (Rust)

Code before:
use std::{
    collections::{hash_map::Entry, HashMap},
    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,
}

fn is_zero(num: &u64) -> bool {
    *num == 0
}

#[derive(Debug, Default)]
pub struct BroadcasterInner {
    state: Mutex<HashMap<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);
        }
    }

Review comment (performance):
```suggestion
#[inline(always)]
fn is_zero(num: &u64) -> bool {
```
may want to do this

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);
Repository: libsql/libsql
File: libsql-server/src/auth/user_auth_strategies/mod.rs (Rust)

Code before:
pub mod disabled;
pub mod http_basic;
pub mod jwt;

pub use disabled::Disabled;
use hashbrown::HashMap;
pub use http_basic::HttpBasic;
pub use jwt::Jwt;

use super::{AuthError, Authenticated};

#[derive(Debug)]
pub struct UserAuthContext {
    pub custom_fields: HashMap<Box<str>, String>, // todo, add aliases
}

impl UserAuthContext {
    pub fn empty() -> UserAuthContext {
        UserAuthContext {
            custom_fields: HashMap::new(),
        }
    }

    pub fn basic(creds: &str) -> UserAuthContext {
        UserAuthContext {
            custom_fields: HashMap::from([("authorization".into(), format!("Basic {creds}"))]),
        }
    }

    pub fn bearer(token: &str) -> UserAuthContext {
        UserAuthContext {
            custom_fields: HashMap::from([("authorization".into(), format!("Bearer {token}"))]),
        }
    }

    pub fn new(scheme: &str, token: &str) -> UserAuthContext {
        UserAuthContext {
            custom_fields: HashMap::from([("authorization".into(), format!("{scheme} {token}"))]),
        }

Review comment (suggestion):
I think we should change the `Box<str>` to `Arc<str>` since these values are well known ahead of time and we can allocate once and avoid having to reallocate a small string on each request via the `.into()` calls.

Code after:
pub mod disabled;
pub mod http_basic;
pub mod jwt;

pub use disabled::Disabled;
use hashbrown::HashMap;
pub use http_basic::HttpBasic;
pub use jwt::Jwt;

use super::{constants::AUTH_HEADER, AuthError, Authenticated};

#[derive(Debug)]
pub struct UserAuthContext {
    custom_fields: HashMap<&'static str, String>,
}

impl UserAuthContext {
    pub fn empty() -> UserAuthContext {
        UserAuthContext {
            custom_fields: HashMap::new(),
        }
    }

    pub fn basic(creds: &str) -> UserAuthContext {
        UserAuthContext {
            custom_fields: HashMap::from([(AUTH_HEADER, format!("Basic {creds}"))]),
        }
    }

    pub fn bearer(token: &str) -> UserAuthContext {
        UserAuthContext {
            custom_fields: HashMap::from([(AUTH_HEADER, format!("Bearer {token}"))]),
        }
    }

    pub fn new(scheme: &str, token: &str) -> UserAuthContext {
        UserAuthContext {
            custom_fields: HashMap::from([(AUTH_HEADER, format!("{scheme} {token}"))]),
        }
Repository: libsql/libsql
File: libsql-server/src/connection/config.rs (Rust)

Code before:
    }
}

impl FromStr for DurabilityMode {
    type Err = ();

    fn from_str(input: &str) -> Result<DurabilityMode, Self::Err> {
        match input {
            "extra" => Ok(DurabilityMode::Extra),
            "strong" => Ok(DurabilityMode::Strong),
            "relaxed" => Ok(DurabilityMode::Relaxed),
            "off" => Ok(DurabilityMode::Off),
            _ => Err(()),
        }
    }
}

impl Display for DurabilityMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let m = match self {
            DurabilityMode::Extra => "extra".to_string(),
            DurabilityMode::Strong => "strong".to_string(),
            DurabilityMode::Relaxed => "relaxed".to_string(),
            DurabilityMode::Off => "off".to_string(),
        };
        write!(f, "{}", m)
    }
}


Review comment (suggestion):
```suggestion
        let m = match self {
            DurabilityMode::Extra => "extra",
            DurabilityMode::Strong => "strong",
            DurabilityMode::Relaxed => "relaxed",
            DurabilityMode::Off => "off",
        };
        write!(f, "{m}")
```


Code after:
        }
    }
}

impl FromStr for DurabilityMode {
    type Err = ();

    fn from_str(input: &str) -> Result<DurabilityMode, Self::Err> {
        match input {
            "extra" => Ok(DurabilityMode::Extra),
            "strong" => Ok(DurabilityMode::Strong),
            "relaxed" => Ok(DurabilityMode::Relaxed),
            "off" => Ok(DurabilityMode::Off),
            _ => Err(()),
        }
    }
}

impl Display for DurabilityMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let m = match self {
            DurabilityMode::Extra => "extra",
            DurabilityMode::Strong => "strong",
            DurabilityMode::Relaxed => "relaxed",
            DurabilityMode::Off => "off",
        };
        write!(f, "{m}")
    }
}

impl From<DurabilityMode> for metadata::DurabilityMode {
    fn from(value: DurabilityMode) -> Self {
        match value {
            DurabilityMode::Relaxed => metadata::DurabilityMode::Relaxed,
            DurabilityMode::Strong => metadata::DurabilityMode::Strong,
            DurabilityMode::Extra => metadata::DurabilityMode::Extra,
            DurabilityMode::Off => metadata::DurabilityMode::Off,
        }
    }
}

impl From<Result<metadata::DurabilityMode, prost::DecodeError>> for DurabilityMode {
    fn from(value: Result<metadata::DurabilityMode, prost::DecodeError>) -> Self {
        match value {
            Ok(mode) => match mode {
                metadata::DurabilityMode::Relaxed => DurabilityMode::Relaxed,
                metadata::DurabilityMode::Strong => DurabilityMode::Strong,
                metadata::DurabilityMode::Extra => DurabilityMode::Extra,
                metadata::DurabilityMode::Off => DurabilityMode::Off,
            },
            Err(_) => DurabilityMode::default(),
Repository: libsql/libsql
File: bindings/c/src/types.rs (Rust)

Code before:

    #[allow(clippy::mut_from_ref)]
    pub fn get_ref_mut(&self) -> &mut libsql::Connection {
        let ptr_mut = self.ptr as *mut libsql_connection;
        &mut unsafe { &mut (*ptr_mut) }.conn
    }
}

#[allow(clippy::from_over_into)]
impl From<&libsql_connection> for libsql_connection_t {
    fn from(value: &libsql_connection) -> Self {
        Self { ptr: value }
    }
}

#[allow(clippy::from_over_into)]
impl From<&mut libsql_connection> for libsql_connection_t {
    fn from(value: &mut libsql_connection) -> Self {
        Self { ptr: value }
    }
}

#[repr(C)]
pub struct replicated {
    pub frame_no: usize,
    pub frames_synced: usize,
}

pub struct stmt {
    pub stmt: libsql::Statement,
    pub params: Vec<libsql::Value>,
}

pub struct libsql_stmt {
    pub stmt: stmt,
}

#[derive(Clone, Debug)]
#[repr(transparent)]
pub struct libsql_stmt_t {
    ptr: *const libsql_stmt,
}

impl libsql_stmt_t {
    pub fn null() -> libsql_stmt_t {
        libsql_stmt_t {
            ptr: std::ptr::null(),
        }
    }

    pub fn is_null(&self) -> bool {

Review comment (suggestion):
Please use types from std::ffi like `std::ffi::c_int`

Code after:

    #[allow(clippy::mut_from_ref)]
    pub fn get_ref_mut(&self) -> &mut libsql::Connection {
        let ptr_mut = self.ptr as *mut libsql_connection;
        &mut unsafe { &mut (*ptr_mut) }.conn
    }
}

#[allow(clippy::from_over_into)]
impl From<&libsql_connection> for libsql_connection_t {
    fn from(value: &libsql_connection) -> Self {
        Self { ptr: value }
    }
}

#[allow(clippy::from_over_into)]
impl From<&mut libsql_connection> for libsql_connection_t {
    fn from(value: &mut libsql_connection) -> Self {
        Self { ptr: value }
    }
}

#[repr(C)]
pub struct replicated {
    pub frame_no: std::ffi::c_int,
    pub frames_synced: std::ffi::c_int,
}

pub struct stmt {
    pub stmt: libsql::Statement,
    pub params: Vec<libsql::Value>,
}

pub struct libsql_stmt {
    pub stmt: stmt,
}

#[derive(Clone, Debug)]
#[repr(transparent)]
pub struct libsql_stmt_t {
    ptr: *const libsql_stmt,
}

impl libsql_stmt_t {
    pub fn null() -> libsql_stmt_t {
        libsql_stmt_t {
            ptr: std::ptr::null(),
        }
    }

    pub fn is_null(&self) -> bool {
Repository: libsql/libsql
File: libsql-server/src/connection/libsql.rs (Rust)

Code before:
        Ok(Self { inner })
    }

    pub fn with_raw<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut rusqlite::Connection) -> R,
    {
        let mut inner = self.inner.lock();
        f(&mut inner.conn)
    }

    pub async fn execute<B: QueryResultBuilder>(
        &self,
        pgm: Program,
        ctx: RequestContext,
        builder: B,
    ) -> Result<(B, Program)> {
        struct Bomb {
            canceled: Arc<AtomicBool>,
            defused: bool,
        }

        impl Drop for Bomb {
            fn drop(&mut self) {
                if !self.defused {
                    tracing::debug!("cancelling request");
                    self.canceled.store(true, Ordering::Relaxed);
                }
            }
        }

        let canceled = {
            let cancelled = self.inner.lock().canceled.clone();
            cancelled.store(false, Ordering::Relaxed);
            cancelled
        };

        let mut bomb = Bomb {
            canceled,
            defused: false,
        };

        PROGRAM_EXEC_COUNT.increment(1);

        check_program_auth(&ctx, &pgm, &self.inner.lock().config_store.get())?;
        let conn = self.inner.clone();
        let ret = BLOCKING_RT
            .spawn_blocking(move || Connection::run(conn, pgm, builder))
            .await
            .unwrap();


Review comment (bug):
can we improve botht he debug logs so they are unique and easier to follow. I think seeing a bunch of these in the logs when turning on debug will be annoying so mayb we consider going to trace? maybe for atleast one of them?

Code after:
        connection_manager.register_connection(&inner, id);

        Ok(Self { inner })
    }

    pub fn with_raw<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut rusqlite::Connection) -> R,
    {
        let mut inner = self.inner.lock();
        f(&mut inner.conn)
    }

    pub async fn execute<B: QueryResultBuilder>(
        &self,
        pgm: Program,
        ctx: RequestContext,
        builder: B,
    ) -> Result<(B, Program)> {
        struct Bomb {
            canceled: Arc<AtomicBool>,
            defused: bool,
        }

        impl Drop for Bomb {
            fn drop(&mut self) {
                if !self.defused {
                    tracing::trace!("cancelling request");
                    self.canceled.store(true, Ordering::Relaxed);
                }
            }
        }

        let canceled = {
            let cancelled = self.inner.lock().canceled.clone();
            cancelled.store(false, Ordering::Relaxed);
            cancelled
        };

        PROGRAM_EXEC_COUNT.increment(1);

        check_program_auth(&ctx, &pgm, &self.inner.lock().config_store.get())?;

        // create the bomb right before spawning the blocking task.
        let mut bomb = Bomb {
            canceled,
            defused: false,
        };
        let conn = self.inner.clone();
        let ret = BLOCKING_RT
            .spawn_blocking(move || Connection::run(conn, pgm, builder))
Repository: libsql/libsql
File: libsql/src/local/connection.rs (Rust)

Code before:
}

pub(crate) struct WalInsertHandle<'a> {
    conn: &'a Connection,
    in_session: RefCell<bool>
}

impl WalInsertHandle<'_> {
    pub fn insert(&self, frame: &[u8]) -> Result<()> {
        assert!(*self.in_session.borrow());
        self.conn.wal_insert_frame(frame)
    }

    pub fn begin(&self) -> Result<()> {
        assert!(!*self.in_session.borrow());
        self.conn.wal_insert_begin()?;
        self.in_session.replace(true);
        Ok(())
    }

    pub fn end(&self) -> Result<()> {
        assert!(*self.in_session.borrow());
        self.conn.wal_insert_end()?;
        self.in_session.replace(false);
        Ok(())
    }
}

impl Drop for WalInsertHandle<'_> {
    fn drop(&mut self) {
        if let Err(err) = self.conn.wal_insert_end() {
            tracing::error!("{:?}", err);
            Err(err).unwrap()
        }
    }
}

impl fmt::Debug for Connection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Connection").finish()
    }
}


Review comment (suggestion):
This will cause `wal_insert_end` to be called twice. Once here and another on `drop`.  Doesn't cause a panic, but might be a problem.

Code after:
}

pub(crate) struct WalInsertHandle<'a> {
    conn: &'a Connection,
    in_session: RefCell<bool>
}

impl WalInsertHandle<'_> {
    pub fn insert(&self, frame: &[u8]) -> Result<()> {
        assert!(*self.in_session.borrow());
        self.conn.wal_insert_frame(frame)
    }

    pub fn begin(&self) -> Result<()> {
        assert!(!*self.in_session.borrow());
        self.conn.wal_insert_begin()?;
        self.in_session.replace(true);
        Ok(())
    }

    pub fn end(&self) -> Result<()> {
        assert!(*self.in_session.borrow());
        self.conn.wal_insert_end()?;
        self.in_session.replace(false);
        Ok(())
    }
}

impl Drop for WalInsertHandle<'_> {
    fn drop(&mut self) {
        if *self.in_session.borrow() {
            if let Err(err) = self.conn.wal_insert_end() {
                tracing::error!("{:?}", err);
                Err(err).unwrap()
            }
        }
    }
}

impl fmt::Debug for Connection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Connection").finish()
    }
}