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

Code before:

#[tonic::async_trait]
impl Proxy for ProxyService {
    type StreamExecStream = Pin<Box<dyn Stream<Item = Result<ExecResp, tonic::Status>> + Send>>;

    async fn stream_exec(
        &self,
        req: tonic::Request<tonic::Streaming<ExecReq>>,
    ) -> Result<tonic::Response<Self::StreamExecStream>, tonic::Status> {
        let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
        let namespace_jwt_key = self
            .namespaces
            .jwt_key(namespace.clone())
            .await
            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))?;
        let auth = if let Some(auth) = &self.auth {
            auth.authenticate_grpc(&req, self.disable_namespaces, namespace_jwt_key)?
        } else {
            Authenticated::from_proxy_grpc_request(&req, self.disable_namespaces)?
        };

        let (connection_maker, _new_frame_notifier) = self
            .namespaces
            .with(namespace, |ns| {
                let connection_maker = ns.db.connection_maker();
                let notifier = ns
                    .db
                    .wal_manager
                    .wrapped()
                    .logger()
                    .new_frame_notifier
                    .subscribe();
                (connection_maker, notifier)
            })
            .await
            .map_err(|e| {
                if let crate::error::Error::NamespaceDoesntExist(_) = e {
                    tonic::Status::failed_precondition(NAMESPACE_DOESNT_EXIST)
                } else {
                    tonic::Status::internal(e.to_string())
                }
            })?;

        let conn = connection_maker.create().await.unwrap();

        let stream = make_proxy_stream(conn, auth, req.into_inner());

        Ok(tonic::Response::new(Box::pin(stream)))
    }

    async fn execute(

Review comment (refactor):
I think we can move this all to a fn and make this more abstracted? You don't need to know the T of the request so you can write a fn that accepts a `req: Request<T>`, that should allow us to unify the ones below.

Code after:
            } else {
                State::Txn.into()
            },
            current_frame_no: last_frame_no,
        });
        Ok(())
    }

    fn into_ret(self) -> Self::Ret {
        self.output.unwrap()
    }
}

// Disconnects all clients that have been idle for more than 30 seconds.
// FIXME: we should also keep a list of recently disconnected clients,
// and if one should arrive with a late message, it should be rejected
// with an error. A similar mechanism is already implemented in hrana-over-http.
pub async fn garbage_collect(clients: &mut HashMap<Uuid, Arc<PrimaryConnection>>) {
    let limit = std::time::Duration::from_secs(30);

    clients.retain(|_, db| db.idle_time() < limit);
    if !clients.is_empty() {
        tracing::trace!("gc: remaining client handles count: {}", clients.len());
    }
}

#[tonic::async_trait]
impl Proxy for ProxyService {
    type StreamExecStream = Pin<Box<dyn Stream<Item = Result<ExecResp, tonic::Status>> + Send>>;

    async fn stream_exec(
        &self,
        mut req: tonic::Request<tonic::Streaming<ExecReq>>,
    ) -> Result<tonic::Response<Self::StreamExecStream>, tonic::Status> {
        let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
        let auth = self.auth(&mut req, namespace.clone()).await?;

        let (connection_maker, _new_frame_notifier) = self
            .namespaces
            .with(namespace, |ns| {
                let connection_maker = ns.db.connection_maker();
                let notifier = ns
                    .db
                    .wal_manager
                    .wrapped()
                    .logger()
                    .new_frame_notifier
                    .subscribe();
                (connection_maker, notifier)
            })
            .await