flush when response exceed limit size

sqld/src/connection/write_proxy.rs

@@ -108,7 +108,6 @@ impl MakeConnection for MakeWriteProxyConn {
     }
 }
 
-#[derive(Debug)]
 pub struct WriteProxyConnection {
     /// Lazily initialized read connection
     read_conn: LibSqlConnection<TransparentMethods>,
@@ -196,8 +195,9 @@ impl WriteProxyConnection {
         let (builder, new_status, new_frame_no) = match res {
             Ok(res) => res,
             Err(e @ (Error::PrimaryStreamDisconnect | Error::PrimaryStreamMisuse)) => {
-                // drop the connection
+                // drop the connection, and reset the state.
                 self.remote_conn.lock().await.take();
+                *status = TxnStatus::Init;
                 return Err(e);
             }
             Err(e) => return Err(e),

sqld/src/rpc/streaming_exec.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use std::task::{ready, Context, Poll};
 
 use futures_core::Stream;
+use prost::Message;
 use rusqlite::types::ValueRef;
 use tokio::sync::mpsc;
 use tonic::{Code, Status};
@@ -50,6 +51,7 @@ struct StreamResponseBuilder {
     request_id: u32,
     sender: mpsc::Sender<ExecResp>,
     current: Option<ProgramResp>,
+    current_size: usize,
 }
 
 impl StreamResponseBuilder {
@@ -59,12 +61,15 @@ impl StreamResponseBuilder {
     }
 
     fn push(&mut self, step: Step) -> Result<(), QueryResultBuilderError> {
-        const MAX_RESPONSE_STEPS: usize = 10;
+        const MAX_RESPONSE_SIZE: usize = bytesize::ByteSize::mb(1).as_u64() as usize;
 
         let current = self.current();
-        current.steps.push(RespStep { step: Some(step) });
+        let step = RespStep { step: Some(step) };
+        let size = step.encoded_len();
+        current.steps.push(step);
+        self.current_size += size;
 
-        if current.steps.len() > MAX_RESPONSE_STEPS {
+        if self.current_size >= MAX_RESPONSE_SIZE {
             self.flush()?;
         }
 
@@ -77,6 +82,7 @@ impl StreamResponseBuilder {
                 request_id: self.request_id,
                 response: Some(exec_resp::Response::ProgramResp(current)),
             };
+            self.current_size = 0;
             self.sender
                 .blocking_send(resp)
                 .map_err(|_| QueryResultBuilderError::Internal(anyhow::anyhow!("stream closed")))?;
@@ -235,6 +241,7 @@ where
                                     request_id,
                                     sender,
                                     current: None,
+                                    current_size: 0,
                                 };
                                 let mut fut = conn.execute_program(pgm, authenticated, builder, None);
                                 loop {

sqld/tests/cluster/mod.rs

@@ -205,40 +205,3 @@ fn sync_many_replica() {
 
     sim.run().unwrap();
 }
-
-#[test]
-fn create_namespace() {
-    let mut sim = Builder::new().build();
-    make_cluster(&mut sim, 0, false);
-
-    sim.client("client", async {
-        let db =
-            Database::open_remote_with_connector("http://foo.primary:8080", "", TurmoilConnector)?;
-        let conn = db.connect()?;
-
-        let Err(e) = conn.execute("create table test (x)", ()).await else {
-            panic!()
-        };
-        assert_snapshot!(e.to_string());
-
-        let client = Client::new();
-        let resp = client
-            .post(
-                "http://foo.primary:9090/v1/namespaces/foo/create",
-                json!({}),
-            )
-            .await?;
-        assert_eq!(resp.status(), 200);
-
-        conn.execute("create table test (x)", ()).await.unwrap();
-        let mut rows = conn.query("select count(*) from test", ()).await.unwrap();
-        assert!(matches!(
-            rows.next().unwrap().unwrap().get_value(0).unwrap(),
-            Value::Integer(0)
-        ));
-
-        Ok(())
-    });
-
-    sim.run().unwrap();
-}

sqld/tests/namespaces/mod.rs

@@ -41,6 +41,43 @@ fn make_primary(sim: &mut Sim, path: PathBuf) {
     });
 }
 
+#[test]
+fn create_namespace() {
+    let mut sim = Builder::new().build();
+    make_cluster(&mut sim, 0, false);
+
+    sim.client("client", async {
+        let db =
+            Database::open_remote_with_connector("http://foo.primary:8080", "", TurmoilConnector)?;
+        let conn = db.connect()?;
+
+        let Err(e) = conn.execute("create table test (x)", ()).await else {
+            panic!()
+        };
+        assert_snapshot!(e.to_string());
+
+        let client = Client::new();
+        let resp = client
+            .post(
+                "http://foo.primary:9090/v1/namespaces/foo/create",
+                json!({}),
+            )
+            .await?;
+        assert_eq!(resp.status(), 200);
+
+        conn.execute("create table test (x)", ()).await.unwrap();
+        let mut rows = conn.query("select count(*) from test", ()).await.unwrap();
+        assert!(matches!(
+            rows.next().unwrap().unwrap().get_value(0).unwrap(),
+            Value::Integer(0)
+        ));
+
+        Ok(())
+    });
+
+    sim.run().unwrap();
+}
+
 #[test]
 fn fork_namespace() {
     let mut sim = Builder::new().build();