refactor: issue request synchronously for MultiReader::commit and MultiWriter::commit

See #6 and #7.

Author: kezhuw

Cargo.toml

@@ -27,6 +27,7 @@ log = "0.4.14"
 static_assertions = "1.1.0"
 hashbrown = "0.12.0"
 hashlink = "0.8.0"
+either = "1.9.0"
 
 [dev-dependencies]
 rand = "0.8.4"

src/client/mod.rs

@@ -5,6 +5,7 @@ use std::future::Future;
 use std::time::Duration;
 
 use const_format::formatcp;
+use either::{Either, Left, Right};
 use thiserror::Error;
 use tokio::sync::{mpsc, watch};
 
@@ -279,15 +280,25 @@ impl Client {
         receiver
     }
 
-    async fn wait<T, F>(result: Result<F>) -> Result<T>
+    async fn wait<T, E, F>(result: std::result::Result<F, E>) -> std::result::Result<T, E>
     where
-        F: Future<Output = Result<T>>, {
+        F: Future<Output = std::result::Result<T, E>>, {
         match result {
             Err(err) => Err(err),
             Ok(future) => future.await,
         }
     }
 
+    async fn resolve<T, E, F>(result: std::result::Result<Either<F, T>, E>) -> std::result::Result<T, E>
+    where
+        F: Future<Output = std::result::Result<T, E>>, {
+        match result {
+            Err(err) => Err(err),
+            Ok(Right(r)) => Ok(r),
+            Ok(Left(future)) => future.await,
+        }
+    }
+
     async fn map_wait<T, U, Fu, Fn>(result: Result<Fu>, f: Fn) -> Result<U>
     where
         Fu: Future<Output = Result<T>>,
@@ -936,8 +947,11 @@ trait MultiBuffer {
     fn op_code() -> OpCode;
 
     fn build_request(&mut self) -> MarshalledRequest {
-        let header = MultiHeader { op: OpCode::Error, done: true, err: -1 };
         let buffer = self.buffer();
+        if buffer.is_empty() {
+            return Default::default();
+        }
+        let header = MultiHeader { op: OpCode::Error, done: true, err: -1 };
         buffer.append_record(&header);
         buffer.finish();
         MarshalledRequest(std::mem::take(buffer))
@@ -1014,23 +1028,32 @@ impl<'a> MultiReader<'a> {
     ///
     /// # Notable behaviors
     /// Individual errors(eg. [Error::NoNode]) are reported individually through [MultiReadResult::Error].
-    pub async fn commit(&mut self) -> Result<Vec<MultiReadResult>> {
-        if self.buf.is_empty() {
-            return Ok(Default::default());
-        }
+    pub fn commit(&mut self) -> impl Future<Output = Result<Vec<MultiReadResult>>> + Send + '_ {
         let request = self.build_request();
+        Client::resolve(self.commit_internally(request))
+    }
+
+    fn commit_internally(
+        &self,
+        request: MarshalledRequest,
+    ) -> Result<Either<impl Future<Output = Result<Vec<MultiReadResult>>> + Send + '_, Vec<MultiReadResult>>> {
+        if request.is_empty() {
+            return Ok(Right(Vec::default()));
+        }
         let receiver = self.client.send_marshalled_request(request);
-        let (body, _) = receiver.await?;
-        let response = record::unmarshal::<Vec<MultiReadResponse>>(&mut body.as_slice())?;
-        let mut results = Vec::with_capacity(response.len());
-        for result in response {
-            match result {
-                MultiReadResponse::Data { data, stat } => results.push(MultiReadResult::Data { data, stat }),
-                MultiReadResponse::Children { children } => results.push(MultiReadResult::Children { children }),
-                MultiReadResponse::Error(err) => results.push(MultiReadResult::Error { err }),
+        Ok(Left(async move {
+            let (body, _) = receiver.await?;
+            let response = record::unmarshal::<Vec<MultiReadResponse>>(&mut body.as_slice())?;
+            let mut results = Vec::with_capacity(response.len());
+            for result in response {
+                match result {
+                    MultiReadResponse::Data { data, stat } => results.push(MultiReadResult::Data { data, stat }),
+                    MultiReadResponse::Children { children } => results.push(MultiReadResult::Children { children }),
+                    MultiReadResponse::Error(err) => results.push(MultiReadResult::Error { err }),
+                }
             }
-        }
-        Ok(results)
+            Ok(results)
+        }))
     }
 
     /// Clears collected operations.
@@ -1184,30 +1207,49 @@ impl<'a> MultiWriter<'a> {
     ///
     /// # Notable errors
     /// * [Error::BadVersion] if check version failed.
-    pub async fn commit(&mut self) -> std::result::Result<Vec<MultiWriteResult>, MultiWriteError> {
-        if self.buf.is_empty() {
-            return Ok(Default::default());
-        }
+    pub fn commit(
+        &mut self,
+    ) -> impl Future<Output = std::result::Result<Vec<MultiWriteResult>, MultiWriteError>> + Send + '_ {
         let request = self.build_request();
+        Client::resolve(self.commit_internally(request))
+    }
+
+    fn commit_internally(
+        &self,
+        request: MarshalledRequest,
+    ) -> std::result::Result<
+        Either<
+            impl Future<Output = std::result::Result<Vec<MultiWriteResult>, MultiWriteError>> + Send + '_,
+            Vec<MultiWriteResult>,
+        >,
+        MultiWriteError,
+    > {
+        if request.is_empty() {
+            return Ok(Right(Vec::default()));
+        }
         let receiver = self.client.send_marshalled_request(request);
-        let (body, _) = receiver.await?;
-        let response = record::unmarshal::<Vec<MultiWriteResponse>>(&mut body.as_slice())?;
-        let failed = response.first().map(|r| matches!(r, MultiWriteResponse::Error(_))).unwrap_or(false);
-        let mut results = if failed { Vec::new() } else { Vec::with_capacity(response.len()) };
-        for (index, result) in response.into_iter().enumerate() {
-            match result {
-                MultiWriteResponse::Check => results.push(MultiWriteResult::Check),
-                MultiWriteResponse::Delete => results.push(MultiWriteResult::Delete),
-                MultiWriteResponse::Create { path, stat } => {
-                    util::strip_root_path(path, self.client.chroot.root())?;
-                    results.push(MultiWriteResult::Create { path: path.to_string(), stat });
-                },
-                MultiWriteResponse::SetData { stat } => results.push(MultiWriteResult::SetData { stat }),
-                MultiWriteResponse::Error(Error::UnexpectedErrorCode(0)) => {},
-                MultiWriteResponse::Error(err) => return Err(MultiWriteError::OperationFailed { index, source: err }),
+        Ok(Left(async move {
+            let (body, _) = receiver.await?;
+            let response = record::unmarshal::<Vec<MultiWriteResponse>>(&mut body.as_slice())?;
+            let failed = response.first().map(|r| matches!(r, MultiWriteResponse::Error(_))).unwrap_or(false);
+            let mut results = if failed { Vec::new() } else { Vec::with_capacity(response.len()) };
+            for (index, result) in response.into_iter().enumerate() {
+                match result {
+                    MultiWriteResponse::Check => results.push(MultiWriteResult::Check),
+                    MultiWriteResponse::Delete => results.push(MultiWriteResult::Delete),
+                    MultiWriteResponse::Create { path, stat } => {
+                        util::strip_root_path(path, self.client.chroot.root())?;
+                        results.push(MultiWriteResult::Create { path: path.to_string(), stat });
+                    },
+                    MultiWriteResponse::SetData { stat } => results.push(MultiWriteResult::SetData { stat }),
+                    MultiWriteResponse::Error(Error::UnexpectedErrorCode(0)) => {},
+                    MultiWriteResponse::Error(err) => {
+                        return Err(MultiWriteError::OperationFailed { index, source: err })
+                    },
+                }
             }
-        }
-        Ok(results)
+            Ok(results)
+        }))
     }
 
     /// Clears collected operations.

src/session/request.rs

@@ -12,7 +12,7 @@ use crate::error::Error;
 use crate::proto::{self, AddWatchMode, ConnectRequest, OpCode, RequestHeader};
 use crate::record::{self, Record, StaticRecord};
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 pub struct MarshalledRequest(pub Vec<u8>);
 
 impl MarshalledRequest {
@@ -82,6 +82,10 @@ impl MarshalledRequest {
         };
         (op_code, watcher_info)
     }
+
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
 }
 
 pub enum Operation {

tests/zookeeper.rs

@@ -215,6 +215,36 @@ async fn test_multi() {
     assert_that!(results).is_empty();
 }
 
+#[tokio::test]
+async fn test_multi_async_order() {
+    let docker = DockerCli::default();
+    let zookeeper = docker.run(zookeeper_image());
+    let zk_port = zookeeper.get_host_port(2181);
+
+    let cluster = format!("127.0.0.1:{}", zk_port);
+    let client = zk::Client::connect(&cluster).await.unwrap();
+
+    let create_options = zk::CreateOptions::new(zk::CreateMode::Persistent, zk::Acl::anyone_all());
+    client.create("/a", "a0".as_bytes(), &create_options).await.unwrap();
+
+    let mut writer = client.new_multi_writer();
+    writer.add_set_data("/a", "a1".as_bytes(), None).unwrap();
+    let write = writer.commit();
+
+    let mut reader = client.new_multi_reader();
+    reader.add_get_data("/a").unwrap();
+    let mut results = reader.commit().await.unwrap();
+    let zk::MultiReadResult::Data { data, stat } = results.remove(0) else { panic!("expect get data result") };
+
+    let mut write_results = write.await.unwrap();
+    let zk::MultiWriteResult::SetData { stat: set_stat } = write_results.remove(0) else {
+        panic!("expect set data result")
+    };
+
+    assert_that!(data).is_equal_to("a1".as_bytes().to_owned());
+    assert_that!(stat).is_equal_to(set_stat);
+}
+
 #[tokio::test]
 async fn test_no_node() {
     let docker = DockerCli::default();