test RemoteConnection

Author: MarinPostma

sqld/src/connection/write_proxy.rs

@@ -2,14 +2,14 @@ use std::path::PathBuf;
 use std::sync::Arc;
 
 use futures_core::future::BoxFuture;
+use futures_core::Stream;
 use parking_lot::Mutex as PMutex;
 use sqld_libsql_bindings::wal_hook::{TransparentMethods, TRANSPARENT_METHODS};
 use tokio::sync::{mpsc, watch, Mutex};
 use tokio_stream::StreamExt;
 use tonic::metadata::BinaryMetadataValue;
 use tonic::transport::Channel;
 use tonic::{Request, Streaming};
-use uuid::Uuid;
 
 use crate::auth::Authenticated;
 use crate::connection::program::{DescribeCol, DescribeParam};
@@ -19,7 +19,7 @@ use crate::query_analysis::TxnStatus;
 use crate::query_result_builder::{QueryBuilderConfig, QueryResultBuilder};
 use crate::replication::FrameNo;
 use crate::rpc::proxy::rpc::proxy_client::ProxyClient;
-use crate::rpc::proxy::rpc::{self, exec_req, exec_resp, DisconnectMessage, ExecReq, ExecResp};
+use crate::rpc::proxy::rpc::{self, exec_req, exec_resp, ExecReq, ExecResp};
 use crate::rpc::NAMESPACE_METADATA_KEY;
 use crate::stats::Stats;
 use crate::{Result, DEFAULT_AUTO_CHECKPOINT};
@@ -102,12 +102,11 @@ impl MakeConnection for MakeWriteProxyConn {
     }
 }
 
-pub struct WriteProxyConnection {
+pub struct WriteProxyConnection<R = Streaming<ExecResp>> {
     /// Lazily initialized read connection
     read_conn: LibSqlConnection<TransparentMethods>,
     write_proxy: ProxyClient<Channel>,
     state: Mutex<TxnStatus>,
-    client_id: Uuid,
     /// FrameNo of the last write performed by this connection on the primary.
     /// any subsequent read on this connection must wait for the replicator to catch up with this
     /// frame_no
@@ -118,7 +117,7 @@ pub struct WriteProxyConnection {
     stats: Arc<Stats>,
     namespace: NamespaceName,
 
-    remote_conn: Mutex<Option<RemoteConnection>>,
+    remote_conn: Mutex<Option<RemoteConnection<R>>>,
 }
 
 impl WriteProxyConnection {
@@ -135,7 +134,6 @@ impl WriteProxyConnection {
             read_conn,
             write_proxy,
             state: Mutex::new(TxnStatus::Init),
-            client_id: Uuid::new_v4(),
             last_write_frame_no: Default::default(),
             applied_frame_no_receiver,
             builder_config,
@@ -234,8 +232,8 @@ impl WriteProxyConnection {
     }
 }
 
-struct RemoteConnection {
-    response_stream: Streaming<ExecResp>,
+struct RemoteConnection<R = Streaming<ExecResp>> {
+    response_stream: R,
     request_sender: mpsc::Sender<ExecReq>,
     current_request_id: u32,
     builder_config: QueryBuilderConfig,
@@ -265,7 +263,12 @@ impl RemoteConnection {
             builder_config,
         })
     }
+}
 
+impl<R> RemoteConnection<R>
+where
+    R: Stream<Item = Result<ExecResp, tonic::Status>> + Unpin,
+{
     /// Perform a request on to the remote peer, and call message_cb for every message received for
     /// that request. message cb should return whether to expect more message for that request.
     async fn make_request(
@@ -458,17 +461,6 @@ impl Connection for WriteProxyConnection {
     }
 }
 
-impl Drop for WriteProxyConnection {
-    fn drop(&mut self) {
-        // best effort attempt to disconnect
-        let mut remote = self.write_proxy.clone();
-        let client_id = self.client_id.to_string();
-        tokio::spawn(async move {
-            let _ = remote.disconnect(DisconnectMessage { client_id }).await;
-        });
-    }
-}
-
 #[cfg(test)]
 pub mod test {
     use arbitrary::{Arbitrary, Unstructured};
@@ -479,7 +471,7 @@ pub mod test {
     use super::*;
     use crate::{
         query_result_builder::{test::test_driver, Column, QueryResultBuilderError},
-        rpc::proxy::rpc::{query_result::RowResult, ExecuteResults},
+        rpc::{proxy::rpc::{query_result::RowResult, ExecuteResults}, streaming_exec::test::random_valid_program_resp},
     };
 
     /// generate an arbitraty rpc value. see build.rs for usage.
@@ -555,4 +547,21 @@ pub mod test {
             },
         );
     }
+
+    #[tokio::test]
+    // in this test we do a roundtrip: generate a random valid program, stream it to
+    // RemoteConnection, and make sure that the remote connection drives the builder with the same
+    // state transitions.
+    async fn validate_random_stream_response() {
+        let (response_stream, validator) = random_valid_program_resp(500, 150);
+        let (request_sender, _request_recver) = mpsc::channel(1);
+        let mut remote = RemoteConnection {
+            response_stream: response_stream.map(Ok),
+            request_sender,
+            current_request_id: 0,
+            builder_config: QueryBuilderConfig::default(),
+        };
+
+        remote.execute(Program::seq(&[]), validator).await.unwrap().0.into_ret();
+    }
 }

sqld/src/http/user/result_builder.rs

@@ -311,15 +311,17 @@ impl QueryResultBuilder for JsonHttpPayloadBuilder {
 
 #[cfg(test)]
 mod test {
-    use crate::query_result_builder::test::random_builder_driver;
+
+    use crate::query_result_builder::test::{random_transition, fsm_builder_driver};
 
     use super::*;
 
     #[test]
     fn test_json_builder() {
         for _ in 0..1000 {
             let builder = JsonHttpPayloadBuilder::new();
-            let ret = random_builder_driver(100, builder).into_ret();
+            let trace = random_transition(100);
+            let ret = fsm_builder_driver(&trace, builder).into_ret();
             println!("{}", std::str::from_utf8(&ret).unwrap());
             // we produce valid json
             serde_json::from_slice::<Vec<serde_json::Value>>(&ret).unwrap();

sqld/src/query_result_builder.rs

@@ -633,7 +633,7 @@ pub mod test {
     #[derive(Debug, PartialEq, Eq, Clone, Copy)]
     #[repr(usize)]
     // do not reorder!
-    enum FsmState {
+    pub enum FsmState {
         Init = 0,
         Finish,
         BeginStep,
@@ -702,11 +702,31 @@ pub mod test {
         }
     }
 
-    pub fn random_builder_driver<B: QueryResultBuilder>(mut max_steps: usize, mut b: B) -> B {
+    pub fn random_transition(mut max_steps: usize) -> Vec<FsmState> {
+        let mut trace = Vec::with_capacity(max_steps);
+        let mut state = Init;
+        trace.push(state);
+        loop {
+            if max_steps > 0 {
+                state = state.rand_transition(false);
+            } else {
+                state = state.toward_finish()
+            }
+
+            trace.push(state);
+            if state == FsmState::Finish {
+                break
+            }
+
+            max_steps = max_steps.saturating_sub(1);
+        }
+        trace
+    }
+
+    pub fn fsm_builder_driver<B: QueryResultBuilder>(trace: &[FsmState], mut b: B) -> B {
         let mut rand_data = [0; 10_000];
         rand_data.try_fill(&mut rand::thread_rng()).unwrap();
         let mut u = Unstructured::new(&rand_data);
-        let mut trace = Vec::new();
 
         #[derive(Arbitrary)]
         pub enum ValueRef<'a> {
@@ -729,9 +749,7 @@ pub mod test {
             }
         }
 
-        let mut state = Init;
-        trace.push(state);
-        loop {
+        for state in trace {
             match state {
                 Init => b.init(&QueryBuilderConfig::default()).unwrap(),
                 BeginStep => b.begin_step().unwrap(),
@@ -758,22 +776,106 @@ pub mod test {
                 }
                 BuilderError => return b,
             }
+        }
 
-            if max_steps > 0 {
-                state = state.rand_transition(false);
-            } else {
-                state = state.toward_finish()
-            }
+        b
+    }
 
-            trace.push(state);
+    /// A Builder that validates a given execution trace
+    pub struct ValidateTraceBuilder {
+        trace: Vec<FsmState>,
+        current: usize,
+    }
 
-            max_steps = max_steps.saturating_sub(1);
+    impl ValidateTraceBuilder {
+        pub fn new(trace: Vec<FsmState>) -> Self {
+            Self { trace, current: 0 }
         }
+    }
 
-        // this can be usefull to help debug the generated test case
-        dbg!(trace);
+    impl QueryResultBuilder for ValidateTraceBuilder {
+        type Ret = ();
 
-        b
+        fn init(&mut self, _config: &QueryBuilderConfig) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::Init);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn begin_step(&mut self) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::BeginStep);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn finish_step(
+            &mut self,
+            _affected_row_count: u64,
+            _last_insert_rowid: Option<i64>,
+        ) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::FinishStep);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn step_error(&mut self, _error: crate::error::Error) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::StepError);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn cols_description<'a>(
+            &mut self,
+            _cols: impl IntoIterator<Item = impl Into<Column<'a>>>,
+        ) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::ColsDescription);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn begin_rows(&mut self) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::BeginRows);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn begin_row(&mut self) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::BeginRow);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn add_row_value(&mut self, _v: ValueRef) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::AddRowValue);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn finish_row(&mut self) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::FinishRow);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn finish_rows(&mut self) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::FinishRows);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn finish(
+            &mut self,
+            _last_frame_no: Option<FrameNo>,
+            _state: TxnStatus,
+        ) -> Result<(), QueryResultBuilderError> {
+            assert_eq!(self.trace[self.current], FsmState::Finish);
+            self.current += 1;
+            Ok(())
+        }
+
+        fn into_ret(self) -> Self::Ret {
+            assert_eq!(self.current, self.trace.len());
+        }
     }
 
     pub struct FsmQueryBuilder {