refactor more

Author: gotocoding-DB

kernel/src/actions/mod.rs

@@ -1712,7 +1712,7 @@ mod tests {
     fn test_commit_info_into_engine_data() {
         let engine = ExprEngine::new();
 
-        let commit_info = CommitInfo::new(0, None, None, None);
+        let commit_info = CommitInfo::new(0, None, None);
         let commit_info_txn_id = commit_info.txn_id.clone();
 
         let engine_data = commit_info.into_engine_data(CommitInfo::to_schema().into(), &engine);

kernel/src/checkpoint/tests.rs

@@ -60,7 +60,11 @@ fn test_deleted_file_retention_timestamp() -> DeltaResult<()> {
 #[test]
 fn test_create_checkpoint_metadata_batch() -> DeltaResult<()> {
     let (store, _) = new_in_memory_store();
-    let engine = DefaultEngine::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
+    let engine = DefaultEngine::new(
+        store.clone(),
+        Arc::new(TokioBackgroundExecutor::new()),
+        None,
+    );
 
     // 1st commit (version 0) - metadata and protocol actions
     // Protocol action does not include the v2Checkpoint reader/writer feature.
@@ -116,7 +120,11 @@ fn test_create_last_checkpoint_data() -> DeltaResult<()> {
     let add_actions_counter = 75;
     let size_in_bytes: i64 = 1024 * 1024; // 1MB
     let (store, _) = new_in_memory_store();
-    let engine = DefaultEngine::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
+    let engine = DefaultEngine::new(
+        store.clone(),
+        Arc::new(TokioBackgroundExecutor::new()),
+        None,
+    );
 
     // Create last checkpoint metadata
     let last_checkpoint_batch = create_last_checkpoint_data(
@@ -277,7 +285,11 @@ fn read_last_checkpoint_file(store: &Arc<InMemory>) -> DeltaResult<Value> {
 #[test]
 fn test_v1_checkpoint_latest_version_by_default() -> DeltaResult<()> {
     let (store, _) = new_in_memory_store();
-    let engine = DefaultEngine::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
+    let engine = DefaultEngine::new(
+        store.clone(),
+        Arc::new(TokioBackgroundExecutor::new()),
+        None,
+    );
 
     // 1st commit: adds `fake_path_1`
     write_commit_to_store(&store, vec![create_add_action("fake_path_1")], 0)?;
@@ -347,7 +359,11 @@ fn test_v1_checkpoint_latest_version_by_default() -> DeltaResult<()> {
 #[test]
 fn test_v1_checkpoint_specific_version() -> DeltaResult<()> {
     let (store, _) = new_in_memory_store();
-    let engine = DefaultEngine::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
+    let engine = DefaultEngine::new(
+        store.clone(),
+        Arc::new(TokioBackgroundExecutor::new()),
+        None,
+    );
 
     // 1st commit (version 0) - metadata and protocol actions
     // Protocol action does not include the v2Checkpoint reader/writer feature.
@@ -409,7 +425,11 @@ fn test_v1_checkpoint_specific_version() -> DeltaResult<()> {
 #[test]
 fn test_finalize_errors_if_checkpoint_data_iterator_is_not_exhausted() -> DeltaResult<()> {
     let (store, _) = new_in_memory_store();
-    let engine = DefaultEngine::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
+    let engine = DefaultEngine::new(
+        store.clone(),
+        Arc::new(TokioBackgroundExecutor::new()),
+        None,
+    );
 
     // 1st commit (version 0) - metadata and protocol actions
     write_commit_to_store(
@@ -451,7 +471,11 @@ fn test_finalize_errors_if_checkpoint_data_iterator_is_not_exhausted() -> DeltaR
 #[test]
 fn test_v2_checkpoint_supported_table() -> DeltaResult<()> {
     let (store, _) = new_in_memory_store();
-    let engine = DefaultEngine::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
+    let engine = DefaultEngine::new(
+        store.clone(),
+        Arc::new(TokioBackgroundExecutor::new()),
+        None,
+    );
 
     // 1st commit: adds `fake_path_2` & removes `fake_path_1`
     write_commit_to_store(
@@ -523,7 +547,11 @@ fn test_v2_checkpoint_supported_table() -> DeltaResult<()> {
 #[test]
 fn test_no_checkpoint_staged_commits() -> DeltaResult<()> {
     let (store, _) = new_in_memory_store();
-    let engine = DefaultEngine::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
+    let engine = DefaultEngine::new(
+        store.clone(),
+        Arc::new(TokioBackgroundExecutor::new()),
+        None,
+    );
 
     // normal commit
     write_commit_to_store(

kernel/src/committer.rs

@@ -212,7 +212,11 @@ mod tests {
     async fn catalog_managed_tables_block_transactions() {
         let storage = Arc::new(InMemory::new());
         let table_root = Url::parse("memory:///").unwrap();
-        let engine = DefaultEngine::new(storage.clone(), Arc::new(TokioBackgroundExecutor::new()));
+        let engine = DefaultEngine::new(
+            storage.clone(),
+            Arc::new(TokioBackgroundExecutor::new()),
+            None,
+        );
 
         let actions = [
             r#"{"commitInfo":{"timestamp":12345678900,"inCommitTimestamp":12345678900}}"#,

kernel/src/engine/default/filesystem.rs

@@ -10,21 +10,102 @@ use url::Url;
 
 use super::UrlExt;
 use crate::engine::default::executor::TaskExecutor;
+use crate::metrics::{MetricEvent, MetricsReporter, Timer};
 use crate::{DeltaResult, Error, FileMeta, FileSlice, StorageHandler};
 
+/// Iterator wrapper that emits metrics when exhausted
+struct ListMetricsIterator<I> {
+    inner: I,
+    reporter: Option<Arc<dyn MetricsReporter>>,
+    timer: Timer,
+    count: u64,
+    metrics_emitted: bool,
+}
+
+impl<I: Iterator<Item = DeltaResult<FileMeta>>> Iterator for ListMetricsIterator<I> {
+    type Item = DeltaResult<FileMeta>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.inner.next() {
+            Some(item) => {
+                if item.is_ok() {
+                    self.count += 1;
+                }
+                Some(item)
+            }
+            None => {
+                if !self.metrics_emitted {
+                    if let Some(ref r) = self.reporter {
+                        r.report(MetricEvent::StorageListCompleted {
+                            duration: self.timer.elapsed(),
+                            num_files: self.count,
+                        });
+                    }
+                    self.metrics_emitted = true;
+                }
+                None
+            }
+        }
+    }
+}
+
+/// Iterator wrapper for read operations that tracks bytes read
+struct ReadMetricsIterator<I> {
+    inner: I,
+    reporter: Option<Arc<dyn MetricsReporter>>,
+    timer: Timer,
+    num_files: u64,
+    bytes_read: u64,
+    metrics_emitted: bool,
+}
+
+impl<I: Iterator<Item = DeltaResult<Bytes>>> Iterator for ReadMetricsIterator<I> {
+    type Item = DeltaResult<Bytes>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.inner.next() {
+            Some(item) => {
+                if let Ok(ref bytes) = item {
+                    self.bytes_read += bytes.len() as u64;
+                }
+                Some(item)
+            }
+            None => {
+                if !self.metrics_emitted {
+                    if let Some(ref r) = self.reporter {
+                        r.report(MetricEvent::StorageReadCompleted {
+                            duration: self.timer.elapsed(),
+                            num_files: self.num_files,
+                            bytes_read: self.bytes_read,
+                        });
+                    }
+                    self.metrics_emitted = true;
+                }
+                None
+            }
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct ObjectStoreStorageHandler<E: TaskExecutor> {
     inner: Arc<DynObjectStore>,
     task_executor: Arc<E>,
+    reporter: Option<Arc<dyn MetricsReporter>>,
     readahead: usize,
 }
 
 impl<E: TaskExecutor> ObjectStoreStorageHandler<E> {
     #[internal_api]
-    pub(crate) fn new(store: Arc<DynObjectStore>, task_executor: Arc<E>) -> Self {
+    pub(crate) fn new(
+        store: Arc<DynObjectStore>,
+        task_executor: Arc<E>,
+        reporter: Option<Arc<dyn MetricsReporter>>,
+    ) -> Self {
         Self {
             inner: store,
             task_executor,
+            reporter,
             readahead: 10,
         }
     }
@@ -106,13 +187,31 @@ impl<E: TaskExecutor> StorageHandler for ObjectStoreStorageHandler<E> {
             }
         });
 
+        let timer = Timer::new();
+        let reporter = self.reporter.clone();
+
         if !has_ordered_listing {
             // This FS doesn't return things in the order we require
             let mut fms: Vec<FileMeta> = receiver.into_iter().try_collect()?;
             fms.sort_unstable();
+
+            let num_files = fms.len() as u64;
+            if let Some(ref r) = reporter {
+                r.report(MetricEvent::StorageListCompleted {
+                    duration: timer.elapsed(),
+                    num_files,
+                });
+            }
+
             Ok(Box::new(fms.into_iter().map(Ok)))
         } else {
-            Ok(Box::new(receiver.into_iter()))
+            Ok(Box::new(ListMetricsIterator {
+                inner: receiver.into_iter(),
+                reporter,
+                timer,
+                count: 0,
+                metrics_emitted: false,
+            }))
         }
     }
 
@@ -127,6 +226,7 @@ impl<E: TaskExecutor> StorageHandler for ObjectStoreStorageHandler<E> {
         files: Vec<FileSlice>,
     ) -> DeltaResult<Box<dyn Iterator<Item = DeltaResult<Bytes>>>> {
         let store = self.inner.clone();
+        let num_files = files.len() as u64;
 
         // This channel will become the output iterator.
         // Because there will already be buffering in the stream, we set the
@@ -173,19 +273,28 @@ impl<E: TaskExecutor> StorageHandler for ObjectStoreStorageHandler<E> {
                 }),
         );
 
-        Ok(Box::new(receiver.into_iter()))
+        Ok(Box::new(ReadMetricsIterator {
+            inner: receiver.into_iter(),
+            reporter: self.reporter.clone(),
+            timer: Timer::new(),
+            num_files,
+            bytes_read: 0,
+            metrics_emitted: false,
+        }))
     }
 
     fn copy_atomic(&self, src: &Url, dest: &Url) -> DeltaResult<()> {
+        let timer = Timer::new();
         let src_path = Path::from_url_path(src.path())?;
         let dest_path = Path::from_url_path(dest.path())?;
         let dest_path_str = dest_path.to_string();
         let store = self.inner.clone();
+        let reporter = self.reporter.clone();
 
         // Read source file then write atomically with PutMode::Create. Note that a GET/PUT is not
         // necessarily atomic, but since the source file is immutable, we aren't exposed to the
         // possiblilty of source file changing while we do the PUT.
-        self.task_executor.block_on(async move {
+        let result = self.task_executor.block_on(async move {
             let data = store.get(&src_path).await?.bytes().await?;
 
             store
@@ -198,7 +307,17 @@ impl<E: TaskExecutor> StorageHandler for ObjectStoreStorageHandler<E> {
                     e => e.into(),
                 })?;
             Ok(())
-        })
+        });
+
+        if result.is_ok() {
+            if let Some(ref r) = reporter {
+                r.report(MetricEvent::StorageCopyCompleted {
+                    duration: timer.elapsed(),
+                });
+            }
+        }
+
+        result
     }
 }
 
@@ -243,7 +362,7 @@ mod tests {
 
         let store = Arc::new(LocalFileSystem::new());
         let executor = Arc::new(TokioBackgroundExecutor::new());
-        let storage = ObjectStoreStorageHandler::new(store, executor);
+        let storage = ObjectStoreStorageHandler::new(store, executor, None);
 
         let mut slices: Vec<FileSlice> = Vec::new();
 
@@ -274,7 +393,7 @@ mod tests {
         store.put(&name, data.clone().into()).await.unwrap();
 
         let table_root = Url::parse("memory:///").expect("valid url");
-        let engine = DefaultEngine::new(store, Arc::new(TokioBackgroundExecutor::new()));
+        let engine = DefaultEngine::new(store, Arc::new(TokioBackgroundExecutor::new()), None);
         let files: Vec<_> = engine
             .storage_handler()
             .list_from(&table_root.join("_delta_log").unwrap().join("0").unwrap())
@@ -304,7 +423,7 @@ mod tests {
 
         let url = Url::from_directory_path(tmp.path()).unwrap();
         let store = Arc::new(LocalFileSystem::new());
-        let engine = DefaultEngine::new(store, Arc::new(TokioBackgroundExecutor::new()));
+        let engine = DefaultEngine::new(store, Arc::new(TokioBackgroundExecutor::new()), None);
         let files = engine
             .storage_handler()
             .list_from(&url.join("_delta_log").unwrap().join("0").unwrap())
@@ -331,7 +450,7 @@ mod tests {
         let tmp = tempfile::tempdir().unwrap();
         let store = Arc::new(LocalFileSystem::new());
         let executor = Arc::new(TokioBackgroundExecutor::new());
-        let handler = ObjectStoreStorageHandler::new(store.clone(), executor);
+        let handler = ObjectStoreStorageHandler::new(store.clone(), executor, None);
 
         // basic
         let data = Bytes::from("test-data");