simplify stream proxy implementation

Author: MarinPostma

Cargo.lock

@@ -71,6 +71,7 @@ rustls-pemfile = "1.0.3"
 rustls = "0.21.7"
 async-stream = "0.3.5"
 libsql = { git = "https://github.com/tursodatabase/libsql.git", rev = "bea8863", optional = true }
+futures-option = "0.2.0"
 
 [dev-dependencies]
 proptest = "1.0.0"

sqld/Cargo.toml

@@ -1,8 +1,10 @@
 use std::collections::HashMap;
+use std::pin::Pin;
 use std::str::FromStr;
 use std::sync::Arc;
 
 use async_lock::{RwLock, RwLockUpgradableReadGuard};
+use futures_core::Stream;
 use rusqlite::types::ValueRef;
 use uuid::Uuid;
 
@@ -15,14 +17,14 @@ use crate::query_result_builder::{
     Column, QueryBuilderConfig, QueryResultBuilder, QueryResultBuilderError,
 };
 use crate::replication::FrameNo;
+use crate::rpc::streaming_exec::make_proxy_stream;
 
 use self::rpc::proxy_server::Proxy;
 use self::rpc::query_result::RowResult;
 use self::rpc::{
     describe_result, Ack, DescribeRequest, DescribeResult, Description, DisconnectMessage, ExecReq,
-    ExecuteResults, QueryResult, ResultRows, Row,
+    ExecuteResults, QueryResult, ResultRows, Row, ExecResp,
 };
-use super::streaming_exec::StreamRequestHandler;
 use super::NAMESPACE_DOESNT_EXIST;
 
 pub mod rpc {
@@ -467,20 +469,18 @@ pub async fn garbage_collect(clients: &mut HashMap<Uuid, Arc<PrimaryConnection>>
 
 #[tonic::async_trait]
 impl Proxy for ProxyService {
-    type StreamExecStream = StreamRequestHandler<tonic::Streaming<ExecReq>>;
+    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> {
-        dbg!();
-        let authenticated = if let Some(auth) = &self.auth {
+        let auth= if let Some(auth) = &self.auth {
             auth.authenticate_grpc(&req, self.disable_namespaces)?
         } else {
             Authenticated::from_proxy_grpc_request(&req, self.disable_namespaces)?
         };
 
-        dbg!();
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
         let (connection_maker, _new_frame_notifier) = self
             .namespaces
@@ -498,13 +498,11 @@ impl Proxy for ProxyService {
                 }
             })?;
 
-        dbg!();
-        let connection = connection_maker.create().await.unwrap();
+        let conn = connection_maker.create().await.unwrap();
 
-        dbg!();
-        let handler = StreamRequestHandler::new(req.into_inner(), connection, authenticated);
+        let stream = make_proxy_stream(conn, auth, req.into_inner());
 
-        Ok(tonic::Response::new(handler))
+        Ok(tonic::Response::new(Box::pin(stream)))
     }
 
     async fn execute(

sqld/src/rpc/proxy.rs

@@ -1,11 +1,12 @@
-use std::pin::Pin;
 use std::sync::Arc;
-use std::task::{ready, Context, Poll};
 
 use futures_core::Stream;
+use futures_option::OptionExt;
 use prost::Message;
 use rusqlite::types::ValueRef;
+use tokio::pin;
 use tokio::sync::mpsc;
+use tokio_stream::StreamExt;
 use tonic::{Code, Status};
 
 use crate::auth::Authenticated;
@@ -17,32 +18,76 @@ use crate::query_result_builder::{
 };
 use crate::replication::FrameNo;
 use crate::rpc::proxy::rpc::exec_req::Request;
-use crate::rpc::proxy::rpc::exec_resp;
+use crate::rpc::proxy::rpc::exec_resp::{self, Response};
 
 use super::proxy::rpc::resp_step::Step;
 use super::proxy::rpc::{self, ExecReq, ExecResp, ProgramResp, RespStep, RowValue};
 
-pin_project_lite::pin_project! {
-    pub struct StreamRequestHandler<S> {
-        #[pin]
-        request_stream: S,
-        connection: Arc<PrimaryConnection>,
-        state: State,
-        authenticated: Authenticated,
-    }
-}
-
-impl<S> StreamRequestHandler<S> {
-    pub fn new(
-        request_stream: S,
-        connection: PrimaryConnection,
-        authenticated: Authenticated,
-    ) -> Self {
-        Self {
-            request_stream,
-            connection: connection.into(),
-            state: State::Idle,
-            authenticated,
+pub fn make_proxy_stream<S>(conn: PrimaryConnection, auth: Authenticated, request_stream: S) -> impl Stream<Item = Result<ExecResp, Status>> 
+where
+    S: Stream<Item = Result<ExecReq, Status>>,
+{
+    async_stream::stream! {
+        let mut current_request_fut = None;
+        let (snd, mut recv) = mpsc::channel(1);
+        let conn = Arc::new(conn);
+        pin!(request_stream);
+
+        loop {
+            tokio::select! {
+                biased;
+                Some(maybe_req) = request_stream.next() => {
+                    match maybe_req {
+                        Err(e) => {
+                            tracing::error!("stream error: {e}");
+                            break
+                        }
+                        Ok(req) => {
+                            let request_id = req.request_id;
+                            match req.request {
+                                Some(Request::Execute(pgm)) => {
+                                    let Ok(pgm) =
+                                        crate::connection::program::Program::try_from(pgm.pgm.unwrap()) else {
+                                            yield Err(Status::new(Code::InvalidArgument, "invalid program"));
+                                            break
+                                        };
+                                    let conn = conn.clone();
+                                    let auth = auth.clone();
+                                    let sender = snd.clone();
+
+                                    let fut = async move {
+                                        let builder = StreamResponseBuilder {
+                                            request_id,
+                                            sender,
+                                            current: None,
+                                            current_size: 0,
+                                        };
+
+                                        let ret = conn.execute_program(pgm, auth, builder, None).await;
+                                        (ret, request_id)
+                                    };
+
+                                    current_request_fut.replace(Box::pin(fut));
+                                }
+                                Some(Request::Describe(_)) => todo!(),
+                                None => {
+                                    yield Err(Status::new(Code::InvalidArgument, "invalid request"));
+                                    break
+                                }
+                            }
+                        }
+                    }
+                },
+                Some(res) = recv.recv() => {
+                    yield Ok(res);
+                },
+                (ret, request_id) = current_request_fut.current(), if current_request_fut.is_some() => {
+                    if let Err(e) = ret {
+                        yield Ok(ExecResp { request_id, response: Some(Response::Error(e.into())) })
+                    }
+                },
+                else => break,
+            }
         }
     }
 }
@@ -199,110 +244,3 @@ impl From<ValueRef<'_>> for RowValue {
         RowValue { value }
     }
 }
-
-enum State {
-    Execute(Pin<Box<dyn Stream<Item = ExecResp> + Send>>),
-    Idle,
-    Fused,
-}
-
-impl<S> Stream for StreamRequestHandler<S>
-where
-    S: Stream<Item = Result<ExecReq, Status>>,
-{
-    type Item = Result<ExecResp, Status>;
-
-    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        let this = self.project();
-
-        // we always poll from the request stream. If a new request arrive, we interupt the current
-        // one, and move to the next.
-        if let Poll::Ready(maybe_req) = this.request_stream.poll_next(cx) {
-            match maybe_req {
-                Some(Err(e)) => {
-                    *this.state = State::Fused;
-                    return Poll::Ready(Some(Err(e)))
-                }
-                Some(Ok(req)) => {
-                    let request_id = req.request_id;
-                    match req.request {
-                        Some(Request::Execute(pgm)) => {
-                            let Ok(pgm) =
-                                crate::connection::program::Program::try_from(pgm.pgm.unwrap()) else {
-                                    *this.state = State::Fused;
-                                    return Poll::Ready(Some(Err(Status::new(Code::InvalidArgument, "invalid program"))));
-                                };
-                            let conn = this.connection.clone();
-                            let authenticated = this.authenticated.clone();
-
-                            let s = async_stream::stream! {
-                                let (sender, mut receiver) = mpsc::channel(1);
-                                let builder = StreamResponseBuilder {
-                                    request_id,
-                                    sender,
-                                    current: None,
-                                    current_size: 0,
-                                };
-                                let mut fut = conn.execute_program(pgm, authenticated, builder, None);
-                                loop {
-                                    tokio::select! {
-                                        res = &mut fut => {
-                                            // drain the receiver
-                                            while let Ok(msg) = receiver.try_recv() {
-                                                yield msg;
-                                            }
-
-                                            if let Err(e) = res {
-                                                yield ExecResp {
-                                                    request_id,
-                                                    response: Some(exec_resp::Response::Error(e.into()))
-                                                }
-                                            }
-                                            break
-                                        }
-                                        msg = receiver.recv() => {
-                                            if let Some(msg) = msg {
-                                                yield msg;
-                                            }
-                                        }
-                                    }
-                                }
-                            };
-                            *this.state = State::Execute(Box::pin(s));
-                        }
-                        Some(Request::Describe(_)) => todo!(),
-                        None => {
-                            *this.state = State::Fused;
-                            return Poll::Ready(Some(Err(Status::new(
-                                            Code::InvalidArgument,
-                                            "invalid ExecReq: missing request",
-                            ))));
-                        }
-                    }
-                }
-                None => {
-                    *this.state = State::Fused;
-                    return Poll::Ready(None)
-                }
-            }
-        }
-
-        match this.state {
-            State::Idle => Poll::Pending,
-            State::Fused => Poll::Ready(None),
-            State::Execute(stream) => {
-                let resp = ready!(stream.as_mut().poll_next(cx));
-                match resp {
-                    Some(resp) => Poll::Ready(Some(Ok(resp))),
-                    None => {
-                        // finished processing this query. Wake up immediately to prepare for the
-                        // next
-                        *this.state = State::Idle;
-                        cx.waker().wake_by_ref();
-                        Poll::Pending
-                    }
-                }
-            }
-        }
-    }
-}