Merge pull request #9062 from Xuanwo/cache-policy

feat: Cache Policy

Author: mergify[bot]

Cargo.lock

@@ -15,7 +15,7 @@ common-arrow = { path = "../arrow" }
 
 anyhow = { workspace = true }
 bincode = { version = "2.0.0-rc.1", features = ["serde", "std", "alloc"] }
-opendal = "0.21"
+opendal = "0.22"
 paste = "1.0.9"
 prost = { workspace = true }
 serde = { workspace = true }

src/common/exception/Cargo.toml

@@ -17,9 +17,11 @@ common-exception = { path = "../exception" }
 anyhow = { workspace = true }
 async-trait = "0.1"
 backon = "0.2"
+bytes = "1"
+futures = "0.3"
 globiter = "0.1"
 once_cell = "1"
-opendal = { version = "0.21", features = [
+opendal = { version = "0.22", features = [
     "layers-tracing",
     "layers-metrics",
     "services-ipfs",

src/common/storage/Cargo.toml

@@ -0,0 +1,105 @@
+//  Copyright 2022 Datafuse 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::sync::Arc;
+
+use async_trait::async_trait;
+use bytes::Bytes;
+use futures::future::BoxFuture;
+use futures::io::Cursor;
+use futures::AsyncReadExt;
+use opendal::layers::CachePolicy;
+use opendal::raw::Accessor;
+use opendal::raw::BytesReader;
+use opendal::raw::RpRead;
+use opendal::Error;
+use opendal::ErrorKind;
+use opendal::OpRead;
+use opendal::OpWrite;
+use opendal::Result;
+
+/// TODO: implement more complex cache logic.
+///
+/// For example:
+///
+/// - Implement a top n heap, and only cache files exist in heap.
+/// - Only cache data file, and ignore snapshot files.
+#[derive(Debug, Default)]
+pub struct FuseCachePolicy {}
+
+impl FuseCachePolicy {
+    pub fn new() -> Self {
+        FuseCachePolicy::default()
+    }
+
+    fn cache_path(&self, path: &str, args: &OpRead) -> String {
+        format!("{path}.cache-{}", args.range().to_header())
+    }
+}
+
+#[async_trait]
+impl CachePolicy for FuseCachePolicy {
+    fn on_read(
+        &self,
+        inner: Arc<dyn Accessor>,
+        cache: Arc<dyn Accessor>,
+        path: &str,
+        args: OpRead,
+    ) -> BoxFuture<'static, Result<(RpRead, BytesReader)>> {
+        let path = path.to_string();
+        let cache_path = self.cache_path(&path, &args);
+        Box::pin(async move {
+            match cache.read(&cache_path, OpRead::default()).await {
+                Ok(v) => Ok(v),
+                Err(err) if err.kind() == ErrorKind::ObjectNotFound => {
+                    let (rp, mut r) = inner.read(&path, args.clone()).await?;
+
+                    let size = rp.clone().into_metadata().content_length();
+                    // If size < 8MiB, we can optimize by buffer in memory.
+                    // TODO: make this configurable.
+                    if size <= 8 * 1024 * 1024 {
+                        let mut bs = Vec::with_capacity(size as usize);
+                        r.read_to_end(&mut bs).await.map_err(|err| {
+                            Error::new(
+                                ErrorKind::Unexpected,
+                                "read from underlying storage service",
+                            )
+                            .set_source(err)
+                        })?;
+                        let bs = Bytes::from(bs);
+
+                        // Ignore errors returned by cache services.
+                        let _ = cache
+                            .write(
+                                &cache_path,
+                                OpWrite::new(size),
+                                Box::new(Cursor::new(bs.clone())),
+                            )
+                            .await;
+                        Ok((rp, Box::new(Cursor::new(bs)) as BytesReader))
+                    } else {
+                        // Ignore errors returned by cache services.
+                        let _ = cache.write(&cache_path, OpWrite::new(size), r).await;
+
+                        match cache.read(&cache_path, OpRead::default()).await {
+                            Ok(v) => Ok(v),
+                            Err(_) => return inner.read(&path, args).await,
+                        }
+                    }
+                }
+                Err(_) => return inner.read(&path, args).await,
+            }
+        })
+    }
+}

src/common/storage/src/cache.rs

@@ -65,3 +65,6 @@ pub use metrics::StorageMetricsLayer;
 
 mod runtime_layer;
 mod utils;
+
+mod cache;
+pub use cache::FuseCachePolicy;

src/common/storage/src/lib.rs

@@ -60,6 +60,30 @@ use crate::StorageRedisConfig;
 
 /// init_operator will init an opendal operator based on storage config.
 pub fn init_operator(cfg: &StorageParams) -> Result<Operator> {
+    let op = init_operator_without_layers(cfg)?;
+
+    let op = op
+        // Add retry
+        .layer(RetryLayer::new(ExponentialBackoff::default().with_jitter()))
+        // Add metrics
+        .layer(MetricsLayer)
+        // Add logging
+        .layer(LoggingLayer::default())
+        // Add tracing
+        .layer(TracingLayer)
+        // NOTE
+        //
+        // Magic happens here. We will add a layer upon original
+        // storage operator so that all underlying storage operations
+        // will send to storage runtime.
+        .layer(RuntimeLayer::new(GlobalIORuntime::instance().inner()));
+
+    Ok(op)
+}
+
+/// init_operator_without_layers will init an opendal operator based
+/// on storage config with any layers.
+pub fn init_operator_without_layers(cfg: &StorageParams) -> Result<Operator> {
     let op = match &cfg {
         StorageParams::Azblob(cfg) => init_azblob_operator(cfg)?,
         StorageParams::Fs(cfg) => init_fs_operator(cfg)?,
@@ -83,22 +107,6 @@ pub fn init_operator(cfg: &StorageParams) -> Result<Operator> {
         }
     };
 
-    let op = op
-        // Add retry
-        .layer(RetryLayer::new(ExponentialBackoff::default().with_jitter()))
-        // Add metrics
-        .layer(MetricsLayer)
-        // Add logging
-        .layer(LoggingLayer)
-        // Add tracing
-        .layer(TracingLayer)
-        // NOTE
-        //
-        // Magic happens here. We will add a layer upon original
-        // storage operator so that all underlying storage operations
-        // will send to storage runtime.
-        .layer(RuntimeLayer::new(GlobalIORuntime::instance().inner()));
-
     Ok(op)
 }
 
@@ -442,7 +450,25 @@ impl CacheOperator {
             return Ok(CacheOperator { op: None });
         }
 
-        let operator = init_operator(&conf.params)?;
+        let operator = init_operator_without_layers(&conf.params)?
+            // Add retry
+            .layer(RetryLayer::new(ExponentialBackoff::default().with_jitter()))
+            // Add metrics
+            .layer(MetricsLayer)
+            // Add logging
+            .layer(
+                LoggingLayer::default()
+                    // Ingore expected errors for logging.
+                    .with_error_level(None),
+            )
+            // Add tracing
+            .layer(TracingLayer)
+            // NOTE
+            //
+            // Magic happens here. We will add a layer upon original
+            // storage operator so that all underlying storage operations
+            // will send to storage runtime.
+            .layer(RuntimeLayer::new(GlobalIORuntime::instance().inner()));
 
         // OpenDAL will send a real request to underlying storage to check whether it works or not.
         // If this check failed, it's highly possible that the users have configured it wrongly.

src/common/storage/src/operator.rs

@@ -31,7 +31,7 @@ crossbeam-channel = "0.5.6"
 csv-core = "0.1.10"
 futures = "0.3.24"
 futures-util = "0.3.24"
-opendal = { version = "0.21", features = ["compress"] }
+opendal = { version = "0.22", features = ["compress"] }
 parking_lot = "0.12.1"
 serde_json = { workspace = true }
 similar-asserts = "1.4.2"

src/query/pipeline/sources/Cargo.toml

@@ -102,7 +102,7 @@ lz4 = "1.24.0"
 metrics = "0.20.1"
 naive-cityhash = "0.2.0"
 once_cell = "1.15.0"
-opendal = { version = "0.21", features = ["layers-tracing", "layers-metrics", "compress"] }
+opendal = { version = "0.22", features = ["layers-tracing", "layers-metrics", "compress"] }
 opensrv-mysql = "0.2.0"
 openssl = { version = "0.10.41", features = ["vendored"] }
 parking_lot = "0.12.1"

src/query/service/Cargo.toml

@@ -27,8 +27,8 @@ use wiremock::ResponseTemplate;
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn test_get_storage_accessor_s3() -> Result<()> {
     let mock_server = MockServer::start().await;
-    Mock::given(method("HEAD"))
-        .and(path("/bucket/.opendal"))
+    Mock::given(method("GET"))
+        .and(path("/bucket"))
         .respond_with(ResponseTemplate::new(404))
         .mount(&mock_server)
         .await;
@@ -39,6 +39,7 @@ async fn test_get_storage_accessor_s3() -> Result<()> {
         region: "us-east-2".to_string(),
         endpoint_url: mock_server.uri(),
         bucket: "bucket".to_string(),
+        disable_credential_loader: true,
         ..Default::default()
     });
 

src/query/service/tests/it/sessions/query_ctx.rs

@@ -19,7 +19,7 @@ common-auth = { path = "../../common/auth" }
 http = "0.2"
 log = "0.4"
 moka = "0.9"
-opendal = "0.21"
+opendal = "0.22"
 reqwest = "0.11"
 serde = { workspace = true }
 serde_json = { workspace = true }