feat(test): add ObjectStore trait simulator support

src/object_store/Cargo.toml

@@ -26,6 +26,7 @@ hyper = { version = "0.14", features = ["tcp", "client"] }
 hyper-rustls = { version = "0.24.2", features = ["webpki-roots"] }
 hyper-tls = "0.5.0"
 itertools = "0.12"
+madsim = "0.2.22"
 opendal = "0.44"
 prometheus = { version = "0.13", features = ["process"] }
 risingwave_common = { workspace = true }

src/object_store/src/object/error.rs

@@ -43,6 +43,9 @@ pub enum ObjectErrorInner {
     #[error("Internal error: {0}")]
     #[construct(skip)]
     Internal(String),
+    #[cfg(madsim)]
+    #[error(transparent)]
+    Sim(#[from] crate::object::sim::SimError),
 }
 
 impl ObjectError {
@@ -80,6 +83,10 @@ impl ObjectError {
             ObjectErrorInner::Mem(e) => {
                 return e.is_object_not_found_error();
             }
+            #[cfg(madsim)]
+            ObjectErrorInner::Sim(e) => {
+                return e.is_object_not_found_error();
+            }
             _ => {}
         };
         false
@@ -108,4 +115,11 @@ impl From<ByteStreamError> for ObjectError {
     }
 }
 
+#[cfg(madsim)]
+impl From<std::io::Error> for ObjectError {
+    fn from(e: std::io::Error) -> Self {
+        ObjectErrorInner::Internal(e.to_string()).into()
+    }
+}
+
 pub type ObjectResult<T> = std::result::Result<T, ObjectError>;

src/object_store/src/object/mod.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#[cfg(madsim)]
+pub mod sim;
 use std::ops::{Range, RangeBounds};
 use std::sync::Arc;
 use std::time::Duration;
@@ -39,6 +41,9 @@ pub use error::*;
 use object_metrics::ObjectStoreMetrics;
 use thiserror_ext::AsReport;
 
+#[cfg(madsim)]
+use self::sim::SimObjectStore;
+
 pub type ObjectStoreRef = Arc<ObjectStoreImpl>;
 pub type ObjectStreamingUploader = MonitoredStreamingUploader;
 
@@ -139,6 +144,8 @@ pub enum ObjectStoreImpl {
     InMem(MonitoredObjectStore<InMemObjectStore>),
     Opendal(MonitoredObjectStore<OpendalObjectStore>),
     S3(MonitoredObjectStore<S3ObjectStore>),
+    #[cfg(madsim)]
+    Sim(MonitoredObjectStore<SimObjectStore>),
 }
 
 macro_rules! dispatch_async {
@@ -165,7 +172,12 @@ macro_rules! object_store_impl_method_body {
                 ObjectStoreImpl::S3(s3) => {
                     $dispatch_macro!(s3, $method_name, path $(, $args)*)
                 },
+                #[cfg(madsim)]
+                ObjectStoreImpl::Sim(in_mem) => {
+                    $dispatch_macro!(in_mem, $method_name, path $(, $args)*)
+                },
             }
+
         }
     };
 }
@@ -179,6 +191,7 @@ macro_rules! object_store_impl_method_body_slice {
     ($object_store:expr, $method_name:ident, $dispatch_macro:ident, $paths:expr $(, $args:expr)*) => {
         {
             let paths_rem = partition_object_store_paths($paths);
+
             match $object_store {
                 ObjectStoreImpl::InMem(in_mem) => {
                     $dispatch_macro!(in_mem, $method_name, &paths_rem $(, $args)*)
@@ -189,6 +202,10 @@ macro_rules! object_store_impl_method_body_slice {
                 ObjectStoreImpl::S3(s3) => {
                     $dispatch_macro!(s3, $method_name, &paths_rem $(, $args)*)
                 },
+                #[cfg(madsim)]
+                ObjectStoreImpl::Sim(in_mem) => {
+                    $dispatch_macro!(in_mem, $method_name, &paths_rem $(, $args)*)
+                },
             }
         }
     };
@@ -247,6 +264,8 @@ impl ObjectStoreImpl {
             ObjectStoreImpl::InMem(store) => store.inner.get_object_prefix(obj_id),
             ObjectStoreImpl::Opendal(store) => store.inner.get_object_prefix(obj_id),
             ObjectStoreImpl::S3(store) => store.inner.get_object_prefix(obj_id),
+            #[cfg(madsim)]
+            ObjectStoreImpl::Sim(store) => store.inner.get_object_prefix(obj_id),
         }
     }
 
@@ -257,6 +276,8 @@ impl ObjectStoreImpl {
                 store.inner.op.info().native_capability().write_can_multi
             }
             ObjectStoreImpl::S3(_) => true,
+            #[cfg(madsim)]
+            ObjectStoreImpl::Sim(_) => true,
         }
     }
 
@@ -265,6 +286,8 @@ impl ObjectStoreImpl {
             ObjectStoreImpl::InMem(store) => store.recv_buffer_size(),
             ObjectStoreImpl::Opendal(store) => store.recv_buffer_size(),
             ObjectStoreImpl::S3(store) => store.recv_buffer_size(),
+            #[cfg(madsim)]
+            ObjectStoreImpl::Sim(store) => store.recv_buffer_size(),
         }
     }
 }
@@ -896,6 +919,10 @@ pub async fn build_remote_object_store(
             }
             ObjectStoreImpl::InMem(InMemObjectStore::shared().monitored(metrics))
         }
+        #[cfg(madsim)]
+        sim if sim.starts_with("sim://") => {
+            ObjectStoreImpl::Sim(SimObjectStore::new(url).monitored(metrics))
+        }
         other => {
             unimplemented!(
                 "{} remote object store only supports s3, minio, gcs, oss, cos, azure blob, hdfs, disk, memory, and memory-shared.",

src/object_store/src/object/sim/client.rs

@@ -0,0 +1,48 @@
+// Copyright 2024 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::fmt::Debug;
+use std::net::SocketAddr;
+
+use madsim::net::{Endpoint, Payload};
+
+use super::service::{Request, Response};
+use crate::object::error::ObjectResult as Result;
+
+#[derive(Debug, Clone)]
+pub struct Client {
+    addr: SocketAddr,
+}
+
+impl Client {
+    pub fn new(addr: SocketAddr) -> Self {
+        Self { addr }
+    }
+
+    pub(crate) async fn send_request(&self, req: Request) -> Result<Response> {
+        let resp = self.send_request_io(req).await?;
+        let resp = *resp
+            .downcast::<Result<Response>>()
+            .expect("failed to downcast");
+        resp
+    }
+
+    async fn send_request_io(&self, req: Request) -> std::io::Result<Payload> {
+        let addr = self.addr;
+        let ep = Endpoint::connect(addr).await?;
+        let (tx, mut rx) = ep.connect1(addr).await?;
+        tx.send(Box::new(req)).await?;
+        rx.recv().await
+    }
+}

src/object_store/src/object/sim/error.rs

@@ -0,0 +1,39 @@
+// Copyright 2024 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum SimError {
+    #[error("NotFound error: {0}")]
+    NotFound(String),
+    #[error("Other error: {0}")]
+    Other(String),
+}
+
+impl SimError {
+    pub fn is_object_not_found_error(&self) -> bool {
+        matches!(self, SimError::NotFound(_))
+    }
+}
+
+impl SimError {
+    pub(crate) fn not_found(msg: impl ToString) -> Self {
+        SimError::NotFound(msg.to_string())
+    }
+
+    pub(crate) fn other(msg: impl ToString) -> Self {
+        SimError::Other(msg.to_string())
+    }
+}