refactor: use usize range instead of BlockLocation to read obj

src/ctl/src/cmd_impl/hummock/sst_dump.rs

@@ -31,7 +31,7 @@ use risingwave_frontend::TableCatalog;
 use risingwave_hummock_sdk::compaction_group::hummock_version_ext::HummockVersionExt;
 use risingwave_hummock_sdk::key::FullKey;
 use risingwave_hummock_sdk::HummockSstableObjectId;
-use risingwave_object_store::object::{BlockLocation, ObjectMetadata, ObjectStoreImpl};
+use risingwave_object_store::object::{ObjectMetadata, ObjectStoreImpl};
 use risingwave_pb::hummock::{Level, SstableInfo};
 use risingwave_rpc_client::MetaClient;
 use risingwave_storage::hummock::value::HummockValue;
@@ -176,20 +176,15 @@ async fn get_meta_offset_from_object(
     obj: &ObjectMetadata,
     obj_store: &ObjectStoreImpl,
 ) -> anyhow::Result<u64> {
-    let meta_offset_loc = BlockLocation {
-        offset: obj.total_size
-            - (
-                // version, magic
-                2 * std::mem::size_of::<u32>() +
-                // footer, checksum
-                2 * std::mem::size_of::<u64>()
-            ),
-        size: std::mem::size_of::<u64>(),
-    };
-    Ok(obj_store
-        .read(&obj.key, Some(meta_offset_loc))
-        .await?
-        .get_u64_le())
+    let start = obj.total_size
+        - (
+            // version, magic
+            2 * std::mem::size_of::<u32>() +
+        // footer, checksum
+        2 * std::mem::size_of::<u64>()
+        );
+    let end = start + std::mem::size_of::<u64>();
+    Ok(obj_store.read(&obj.key, start..end).await?.get_u64_le())
 }
 
 pub async fn sst_dump_via_sstable_store(
@@ -281,11 +276,8 @@ async fn print_block(
 
     // Retrieve encoded block data in bytes
     let store = sstable_store.store();
-    let block_loc = BlockLocation {
-        offset: block_meta.offset as usize,
-        size: block_meta.len as usize,
-    };
-    let block_data = store.read(&data_path, Some(block_loc)).await?;
+    let range = block_meta.offset as usize..block_meta.offset as usize + block_meta.len as usize;
+    let block_data = store.read(&data_path, range).await?;
 
     // Retrieve checksum and compression algorithm used from the encoded block data
     let len = block_data.len();

src/meta/src/hummock/manager/checkpoint.rs

@@ -40,7 +40,7 @@ impl HummockManager {
         use prost::Message;
         let data = match self
             .object_store
-            .read(&self.version_checkpoint_path, None)
+            .read(&self.version_checkpoint_path, ..)
             .await
         {
             Ok(data) => data,

src/meta/src/hummock/manager/mod.rs

@@ -2843,7 +2843,7 @@ async fn write_exclusive_cluster_id(
     const CLUSTER_ID_NAME: &str = "0";
     let cluster_id_dir = format!("{}/{}/", state_store_dir, CLUSTER_ID_DIR);
     let cluster_id_full_path = format!("{}{}", cluster_id_dir, CLUSTER_ID_NAME);
-    match object_store.read(&cluster_id_full_path, None).await {
+    match object_store.read(&cluster_id_full_path, ..).await {
         Ok(cluster_id) => Err(ObjectError::internal(format!(
             "Data directory is already used by another cluster with id {:?}, path {}.",
             String::from_utf8(cluster_id.to_vec()).unwrap(),

src/object_store/src/lib.rs

@@ -18,5 +18,6 @@
 #![feature(lint_reasons)]
 #![feature(error_generic_member_access)]
 #![feature(let_chains)]
+#![feature(bound_map)]
 
 pub mod object;

src/object_store/src/object/mem.rs

@@ -14,22 +14,22 @@
 
 use std::collections::{HashMap, VecDeque};
 use std::io::Cursor;
+use std::ops::RangeBounds;
 use std::pin::Pin;
 use std::sync::{Arc, LazyLock};
 use std::task::{Context, Poll};
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use bytes::{BufMut, Bytes, BytesMut};
 use fail::fail_point;
-use futures::future::try_join_all;
 use futures::Stream;
 use itertools::Itertools;
 use thiserror::Error;
 use tokio::io::AsyncRead;
 use tokio::sync::Mutex;
 
 use super::{
-    BlockLocation, BoxedStreamingUploader, ObjectError, ObjectMetadata, ObjectResult, ObjectStore,
+    BoxedStreamingUploader, ObjectError, ObjectMetadata, ObjectResult, ObjectStore,
     StreamingUploader,
 };
 use crate::object::ObjectMetadataIter;
@@ -130,23 +130,15 @@ impl ObjectStore for InMemObjectStore {
         }))
     }
 
-    async fn read(&self, path: &str, block: Option<BlockLocation>) -> ObjectResult<Bytes> {
+    async fn read(
+        &self,
+        path: &str,
+        range: impl RangeBounds<usize> + Clone + Send + Sync + 'static,
+    ) -> ObjectResult<Bytes> {
         fail_point!("mem_read_err", |_| Err(ObjectError::internal(
             "mem read error"
         )));
-        if let Some(loc) = block {
-            self.get_object(path, |obj| find_block(obj, loc)).await?
-        } else {
-            self.get_object(path, |obj| Ok(obj.clone())).await?
-        }
-    }
-
-    async fn readv(&self, path: &str, block_locs: &[BlockLocation]) -> ObjectResult<Vec<Bytes>> {
-        let futures = block_locs
-            .iter()
-            .map(|block_loc| self.read(path, Some(*block_loc)))
-            .collect_vec();
-        try_join_all(futures).await
+        self.get_object(path, range).await
     }
 
     /// Returns a stream reading the object specified in `path`. If given, the stream starts at the
@@ -160,23 +152,10 @@ impl ObjectStore for InMemObjectStore {
         fail_point!("mem_streaming_read_err", |_| Err(ObjectError::internal(
             "mem streaming read error"
         )));
-
-        let bytes = if let Some(pos) = start_pos {
-            self.get_object(path, |obj| {
-                find_block(
-                    obj,
-                    BlockLocation {
-                        offset: pos,
-                        size: obj.len() - pos,
-                    },
-                )
-            })
-            .await?
-        } else {
-            self.get_object(path, |obj| Ok(obj.clone())).await?
-        };
-
-        Ok(Box::new(Cursor::new(bytes?)))
+        let bytes = self
+            .get_object(path, start_pos.unwrap_or_default()..)
+            .await?;
+        Ok(Box::new(Cursor::new(bytes)))
     }
 
     async fn metadata(&self, path: &str) -> ObjectResult<ObjectMetadata> {
@@ -254,25 +233,34 @@ impl InMemObjectStore {
         *SHARED.lock() = InMemObjectStore::new();
     }
 
-    async fn get_object<R, F>(&self, path: &str, f: F) -> ObjectResult<R>
-    where
-        F: Fn(&Bytes) -> R,
-    {
-        self.objects
-            .lock()
-            .await
+    async fn get_object(
+        &self,
+        path: &str,
+        range: impl RangeBounds<usize> + Clone + Send + Sync + 'static,
+    ) -> ObjectResult<Bytes> {
+        let objects = self.objects.lock().await;
+
+        let obj = objects
             .get(path)
             .map(|(_, obj)| obj)
-            .ok_or_else(|| Error::not_found(format!("no object at path '{}'", path)).into())
-            .map(f)
-    }
-}
+            .ok_or_else(|| Error::not_found(format!("no object at path '{}'", path)))?;
+
+        let start = match range.start_bound() {
+            std::ops::Bound::Included(v) => *v,
+            std::ops::Bound::Excluded(v) => *v - 1,
+            std::ops::Bound::Unbounded => 0,
+        };
+        let end = match range.end_bound() {
+            std::ops::Bound::Included(v) => *v + 1,
+            std::ops::Bound::Excluded(v) => *v,
+            std::ops::Bound::Unbounded => obj.len(),
+        };
+
+        if end > obj.len() {
+            return Err(Error::other("bad block offset and size").into());
+        }
 
-fn find_block(obj: &Bytes, block: BlockLocation) -> ObjectResult<Bytes> {
-    if block.offset + block.size > obj.len() {
-        Err(Error::other("bad block offset and size").into())
-    } else {
-        Ok(obj.slice(block.offset..(block.offset + block.size)))
+        Ok(obj.slice(start..end))
     }
 }
 
@@ -326,29 +314,19 @@ mod tests {
         s3.upload("/abc", block).await.unwrap();
 
         // No such object.
-        let err = s3
-            .read("/ab", Some(BlockLocation { offset: 0, size: 3 }))
-            .await
-            .unwrap_err();
+        let err = s3.read("/ab", 0..3).await.unwrap_err();
         assert!(err.is_object_not_found_error());
 
-        let bytes = s3
-            .read("/abc", Some(BlockLocation { offset: 4, size: 2 }))
-            .await
-            .unwrap();
+        let bytes = s3.read("/abc", 4..6).await.unwrap();
         assert_eq!(String::from_utf8(bytes.to_vec()).unwrap(), "56".to_string());
 
         // Overflow.
-        s3.read("/abc", Some(BlockLocation { offset: 4, size: 4 }))
-            .await
-            .unwrap_err();
+        s3.read("/abc", 4..8).await.unwrap_err();
 
         s3.delete("/abc").await.unwrap();
 
         // No such object.
-        s3.read("/abc", Some(BlockLocation { offset: 0, size: 3 }))
-            .await
-            .unwrap_err();
+        s3.read("/abc", 0..3).await.unwrap_err();
     }
 
     #[tokio::test]
@@ -365,14 +343,11 @@ mod tests {
         uploader.finish().await.unwrap();
 
         // Read whole object.
-        let read_obj = store.read("/abc", None).await.unwrap();
+        let read_obj = store.read("/abc", ..).await.unwrap();
         assert!(read_obj.eq(&obj));
 
         // Read part of the object.
-        let read_obj = store
-            .read("/abc", Some(BlockLocation { offset: 4, size: 2 }))
-            .await
-            .unwrap();
+        let read_obj = store.read("/abc", 4..6).await.unwrap();
         assert_eq!(
             String::from_utf8(read_obj.to_vec()).unwrap(),
             "56".to_string()