feat: enable caching when using object store

src/datanode/src/datanode.rs

@@ -25,6 +25,8 @@ use crate::error::Result;
 use crate::instance::{Instance, InstanceRef};
 use crate::server::Services;
 
+pub const DEFAULT_OBJECT_STORE_CACHE_SIZE: ReadableSize = ReadableSize(1024);
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "type")]
 pub enum ObjectStoreConfig {
@@ -48,6 +50,8 @@ pub struct S3Config {
     pub secret_access_key: String,
     pub endpoint: Option<String>,
     pub region: Option<String>,
+    pub cache_path: Option<String>,
+    pub cache_capacity: Option<ReadableSize>,
 }
 
 #[derive(Debug, Clone, Serialize, Default, Deserialize)]
@@ -58,6 +62,8 @@ pub struct OssConfig {
     pub access_key_id: String,
     pub access_key_secret: String,
     pub endpoint: String,
+    pub cache_path: Option<String>,
+    pub cache_capacity: Option<ReadableSize>,
 }
 
 impl Default for ObjectStoreConfig {

src/datanode/src/instance.rs

@@ -19,6 +19,7 @@ use std::{fs, path};
 use backon::ExponentialBackoff;
 use catalog::remote::MetaKvBackend;
 use catalog::{CatalogManager, CatalogManagerRef, RegisterTableRequest};
+use common_base::readable_size::ReadableSize;
 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, MIN_USER_TABLE_ID};
 use common_grpc::channel_manager::{ChannelConfig, ChannelManager};
 use common_telemetry::logging::info;
@@ -28,7 +29,8 @@ use meta_client::client::{MetaClient, MetaClientBuilder};
 use meta_client::MetaClientOpts;
 use mito::config::EngineConfig as TableEngineConfig;
 use mito::engine::MitoEngine;
-use object_store::layers::{LoggingLayer, MetricsLayer, RetryLayer, TracingLayer};
+use object_store::cache_policy::LruCachePolicy;
+use object_store::layers::{CacheLayer, LoggingLayer, MetricsLayer, RetryLayer, TracingLayer};
 use object_store::services::fs::Builder as FsBuilder;
 use object_store::services::oss::Builder as OSSBuilder;
 use object_store::services::s3::Builder as S3Builder;
@@ -42,7 +44,9 @@ use table::table::numbers::NumbersTable;
 use table::table::TableIdProviderRef;
 use table::Table;
 
-use crate::datanode::{DatanodeOptions, ObjectStoreConfig, WalConfig};
+use crate::datanode::{
+    DatanodeOptions, ObjectStoreConfig, WalConfig, DEFAULT_OBJECT_STORE_CACHE_SIZE,
+};
 use crate::error::{
     self, CatalogSnafu, MetaClientInitSnafu, MissingMetasrvOptsSnafu, MissingNodeIdSnafu,
     NewCatalogSnafu, OpenLogStoreSnafu, Result,
@@ -240,7 +244,44 @@ pub(crate) async fn new_oss_object_store(store_config: &ObjectStoreConfig) -> Re
         config: store_config.clone(),
     })?;
 
-    Ok(ObjectStore::new(accessor))
+    create_object_store_with_cache(ObjectStore::new(accessor), store_config)
+}
+
+fn create_object_store_with_cache(
+    object_store: ObjectStore,
+    store_config: &ObjectStoreConfig,
+) -> Result<ObjectStore> {
+    let (cache_path, cache_capacity) = match store_config {
+        ObjectStoreConfig::S3(s3_config) => {
+            let path = s3_config.cache_path.as_ref();
+            let capacity = s3_config
+                .cache_capacity
+                .unwrap_or(DEFAULT_OBJECT_STORE_CACHE_SIZE);
+            (path, capacity)
+        }
+        ObjectStoreConfig::Oss(oss_config) => {
+            let path = oss_config.cache_path.as_ref();
+            let capacity = oss_config
+                .cache_capacity
+                .unwrap_or(DEFAULT_OBJECT_STORE_CACHE_SIZE);
+            (path, capacity)
+        }
+        _ => (None, ReadableSize(0)),
+    };
+
+    if let Some(path) = cache_path {
+        let cache_store =
+            ObjectStore::new(FsBuilder::default().root(path).build().with_context(|_| {
+                error::InitBackendSnafu {
+                    config: store_config.clone(),
+                }
+            })?);
+        let policy = LruCachePolicy::new(cache_capacity.0 as usize);
+        let cache_layer = CacheLayer::new(cache_store).with_policy(policy);
+        Ok(object_store.layer(cache_layer))
+    } else {
+        Ok(object_store)
+    }
 }
 
 pub(crate) async fn new_s3_object_store(store_config: &ObjectStoreConfig) -> Result<ObjectStore> {
@@ -273,7 +314,7 @@ pub(crate) async fn new_s3_object_store(store_config: &ObjectStoreConfig) -> Res
         config: store_config.clone(),
     })?;
 
-    Ok(ObjectStore::new(accessor))
+    create_object_store_with_cache(ObjectStore::new(accessor), store_config)
 }
 
 pub(crate) async fn new_fs_object_store(store_config: &ObjectStoreConfig) -> Result<ObjectStore> {

src/object-store/Cargo.toml

@@ -5,6 +5,8 @@ edition.workspace = true
 license.workspace = true
 
 [dependencies]
+lru = "0.9"
+async-trait = "0.1"
 futures = { version = "0.3" }
 opendal = { version = "0.25.1", features = [
     "layers-tracing",

src/object-store/src/cache_policy.rs

@@ -0,0 +1,125 @@
+// Copyright 2023 Greptime Team
+//
+// 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::num::NonZeroUsize;
+use std::ops::DerefMut;
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use futures::future::BoxFuture;
+use lru::LruCache;
+use opendal::layers::CachePolicy;
+use opendal::raw::output::Reader;
+use opendal::raw::{Accessor, RpDelete, RpRead};
+use opendal::{ErrorKind, OpDelete, OpRead, OpWrite, Result};
+use tokio::sync::Mutex;
+
+struct LruValue;
+
+#[derive(Debug)]
+pub struct LruCachePolicy {
+    lru_cache: Arc<Mutex<LruCache<String, LruValue>>>,
+}
+
+impl LruCachePolicy {
+    pub fn new(capacity: usize) -> Self {
+        Self {
+            lru_cache: Arc::new(Mutex::new(LruCache::new(
+                NonZeroUsize::new(capacity).unwrap(),
+            ))),
+        }
+    }
+
+    fn cache_path(&self, path: &str, args: &OpRead) -> String {
+        format!("{}.cache-{}", path, args.range().to_header())
+    }
+}
+
+#[async_trait]
+impl CachePolicy for LruCachePolicy {
+    fn on_read(
+        &self,
+        inner: Arc<dyn Accessor>,
+        cache: Arc<dyn Accessor>,
+        path: &str,
+        args: OpRead,
+    ) -> BoxFuture<'static, Result<(RpRead, Reader)>> {
+        let path = path.to_string();
+        let cache_path = self.cache_path(&path, &args);
+        let lru_cache = self.lru_cache.clone();
+        Box::pin(async move {
+            match cache.read(&cache_path, OpRead::default()).await {
+                Ok(v) => {
+                    // update lru when cache hit
+                    let mut lru_cache = lru_cache.lock().await;
+                    lru_cache.get_or_insert(cache_path.clone(), || LruValue {});
+                    Ok(v)
+                }
+                Err(err) if err.kind() == ErrorKind::ObjectNotFound => {
+                    let (rp, reader) = inner.read(&path, args.clone()).await?;
+                    let size = rp.clone().into_metadata().content_length();
+                    let _ = cache
+                        .write(&cache_path, OpWrite::new(size), Box::new(reader))
+                        .await?;
+                    match cache.read(&cache_path, OpRead::default()).await {
+                        Ok(v) => {
+                            let r = {
+                                // push new cache file name to lru
+                                let mut lru_cache = lru_cache.lock().await;
+                                lru_cache.push(cache_path.clone(), LruValue {})
+                            };
+                            // delete the evicted cache file
+                            if let Some((k, _v)) = r {
+                                let _ = cache.delete(&k, OpDelete::new()).await;
+                            }
+                            Ok(v)
+                        }
+                        Err(_) => return inner.read(&path, args).await,
+                    }
+                }
+                Err(_) => return inner.read(&path, args).await,
+            }
+        })
+    }
+
+    fn on_delete(
+        &self,
+        inner: Arc<dyn Accessor>,
+        cache: Arc<dyn Accessor>,
+        path: &str,
+        args: OpDelete,
+    ) -> BoxFuture<'static, Result<RpDelete>> {
+        let path = path.to_string();
+        let lru_cache = self.lru_cache.clone();
+        Box::pin(async move {
+            let cache_files: Vec<String> = {
+                let mut guard = lru_cache.lock().await;
+                let lru = guard.deref_mut();
+                let cache_files = lru
+                    .iter()
+                    .filter(|(k, _v)| k.starts_with(format!("{path}.cache-").as_str()))
+                    .map(|(k, _v)| k.clone())
+                    .collect::<Vec<_>>();
+                for k in &cache_files {
+                    lru.pop(k);
+                }
+                cache_files
+            };
+            for file in cache_files {
+                let _ = cache.delete(&file, OpDelete::new()).await;
+            }
+            return inner.delete(&path, args).await;
+        })
+    }
+}

src/object-store/src/lib.rs

@@ -17,5 +17,6 @@ pub use opendal::{
     Operator as ObjectStore, Result,
 };
 pub mod backend;
+pub mod cache_policy;
 pub mod test_util;
 pub mod util;