Merge branch 'main' into null2

Author: mergify[bot]

Cargo.lock

@@ -65,11 +65,9 @@ pub static GLOBAL_MEM_STAT: MemStat = MemStat::empty();
 #[thread_local]
 static mut TRACKER: ThreadTracker = ThreadTracker::empty();
 
-/// Flag indicating an exceeding limit panic is happening and allocating memory by panic handler is allowed.
-///
-/// Flag will be reset when `panic!()` returns.
+/// Whether to allow unlimited memory. Alloc memory will not panic if it is true.
 #[thread_local]
-static PANICKING: AtomicBool = AtomicBool::new(false);
+static UNLIMITED_FLAG: AtomicBool = AtomicBool::new(false);
 
 static MEM_STAT_BUFFER_SIZE: i64 = 4 * 1024 * 1024;
 
@@ -85,24 +83,27 @@ impl<'a> Drop for Entered<'a> {
     }
 }
 
-/// A guard that resets the `PANICKING` flag when dropped.
-pub(crate) struct Panicking;
+/// A guard that resets the `UNLIMITED_FLAG` flag when dropped.
+pub struct UnlimitedMemGuard {
+    saved: bool,
+}
 
-impl Panicking {
-    #[must_use]
-    pub(crate) fn enter_panicking() -> Self {
-        PANICKING.store(true, Ordering::Relaxed);
-        Self
+impl UnlimitedMemGuard {
+    #[allow(unused)]
+    pub(crate) fn enter_unlimited() -> Self {
+        let saved = UNLIMITED_FLAG.load(Ordering::Relaxed);
+        UNLIMITED_FLAG.store(true, Ordering::Relaxed);
+        Self { saved }
     }
 
-    pub(crate) fn is_panicking() -> bool {
-        PANICKING.load(Ordering::Relaxed)
+    pub(crate) fn is_unlimited() -> bool {
+        UNLIMITED_FLAG.load(Ordering::Relaxed)
     }
 }
 
-impl Drop for Panicking {
+impl Drop for UnlimitedMemGuard {
     fn drop(&mut self) {
-        PANICKING.store(false, Ordering::Relaxed);
+        UNLIMITED_FLAG.store(self.saved, Ordering::Relaxed);
     }
 }
 
@@ -205,16 +206,10 @@ impl ThreadTracker {
         let res = tracker.flush();
 
         if let Err(out_of_limit) = res {
-            // NOTE: `PANICKING` only allows allocation inside the following `panic!()`.
-            //       If a `Drop` is called when unwinding, the `Drop` may panic again if it allocates memory over the limit.
-            if Panicking::is_panicking() {
-                return;
+            // https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=03d21a15e52c7c0356fca04ece283cf9
+            if !std::thread::panicking() && !UnlimitedMemGuard::is_unlimited() {
+                panic!("{:?}", out_of_limit);
             }
-
-            // Reset PANICKING when dropped.
-            let _p = Panicking::enter_panicking();
-
-            panic!("{:?}", out_of_limit);
         }
     }
 
@@ -420,7 +415,6 @@ impl<T: Future> Future for TrackedFuture<T> {
 
 #[cfg(test)]
 mod tests {
-
     mod async_thread_tracker {
         use std::future::Future;
         use std::pin::Pin;

src/common/base/src/base/runtime_tracker.rs

@@ -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

@@ -47,7 +47,8 @@ where R: BufferRead
 {
     fn read_date_text(&mut self, tz: &Tz) -> Result<NaiveDate> {
         // TODO support YYYYMMDD format
-        self.read_timestamp_text(tz).map(|dt| dt.naive_utc().date())
+        self.read_timestamp_text(tz)
+            .map(|dt| dt.naive_local().date())
     }
 
     fn read_timestamp_text(&mut self, tz: &Tz) -> Result<DateTime<Tz>> {

src/common/io/src/buffer/buffer_read_datetime_ext.rs

@@ -75,7 +75,8 @@ where T: AsRef<[u8]>
 {
     fn read_date_text(&mut self, tz: &Tz) -> Result<NaiveDate> {
         // TODO support YYYYMMDD format
-        self.read_timestamp_text(tz).map(|dt| dt.naive_utc().date())
+        self.read_timestamp_text(tz)
+            .map(|dt| dt.naive_local().date())
     }
 
     fn read_timestamp_text(&mut self, tz: &Tz) -> Result<DateTime<Tz>> {

src/common/io/src/cursor_ext/cursor_read_datetime_ext.rs

@@ -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;