feat: Add storage stats for IoStatistics

kewang1024 · Jan 24, 2025 · aed3c40 · aed3c40
1 parent 419de77
commit aed3c40
Show file tree

Hide file tree

Showing 10 changed files with 28 additions and 7 deletions.
diff --git a/velox/common/file/FileSystems.cpp b/velox/common/file/FileSystems.cpp
@@ -115,7 +115,8 @@ class LocalFileSystem : public FileSystem {
 
   std::unique_ptr<ReadFile> openFileForRead(
       std::string_view path,
-      const FileOptions& options) override {
+      const FileOptions& options,
+      io::IoStatistics* ioStats) override {
     return std::make_unique<LocalReadFile>(
         extractPath(path), executor_.get(), options.bufferIo);
   }

diff --git a/velox/common/file/FileSystems.h b/velox/common/file/FileSystems.h
@@ -16,6 +16,7 @@
 #pragma once
 
 #include "velox/common/base/Exceptions.h"
+#include "velox/common/io/IoStatistics.h"
 #include "velox/common/memory/MemoryPool.h"
 
 #include <functional>
@@ -103,7 +104,8 @@ class FileSystem {
   /// Returns a ReadFile handle for a given file path
   virtual std::unique_ptr<ReadFile> openFileForRead(
       std::string_view path,
-      const FileOptions& options = {}) = 0;
+      const FileOptions& options = {},
+      io::IoStatistics* ioStats = {}) = 0;
 
   /// Returns a WriteFile handle for a given file path
   virtual std::unique_ptr<WriteFile> openFileForWrite(

diff --git a/velox/common/file/tests/FaultyFileSystem.h b/velox/common/file/tests/FaultyFileSystem.h
@@ -54,7 +54,8 @@ class FaultyFileSystem : public FileSystem {
 
   std::unique_ptr<ReadFile> openFileForRead(
       std::string_view path,
-      const FileOptions& options = {}) override;
+      const FileOptions& options = {},
+      io::IoStatistics* ioStats = {}) override;
 
   std::unique_ptr<WriteFile> openFileForWrite(
       std::string_view path,

diff --git a/velox/common/io/IoStatistics.cpp b/velox/common/io/IoStatistics.cpp
@@ -108,6 +108,12 @@ IoStatistics::operationStats() const {
   return operationStats_;
 }
 
+std::unordered_map<std::string, RuntimeCounter> IoStatistics::storageStats()
+    const {
+  std::lock_guard<std::mutex> lock{storageStatsMutex_};
+  return storageStats_;
+}
+
 void IoStatistics::merge(const IoStatistics& other) {
   rawBytesRead_ += other.rawBytesRead_;
   rawBytesWritten_ += other.rawBytesWritten_;

diff --git a/velox/common/io/IoStatistics.h b/velox/common/io/IoStatistics.h
@@ -23,6 +23,7 @@
 #include <unordered_map>
 
 #include <folly/dynamic.h>
+#include "velox/common/base/RuntimeMetrics.h"
 
 namespace facebook::velox::io {
 
@@ -140,6 +141,7 @@ class IoStatistics {
       const uint64_t partialThrottleCount = 0);
 
   std::unordered_map<std::string, OperationCounters> operationStats() const;
+  std::unordered_map<std::string, RuntimeCounter> storageStats() const;
 
   void merge(const IoStatistics& other);
 
@@ -172,7 +174,9 @@ class IoStatistics {
   IoCounter queryThreadIoLatency_;
 
   std::unordered_map<std::string, OperationCounters> operationStats_;
+  std::unordered_map<std::string, RuntimeCounter> storageStats_;
   mutable std::mutex operationStatsMutex_;
+  mutable std::mutex storageStatsMutex_;
 };
 
 } // namespace facebook::velox::io
diff --git a/velox/connectors/hive/HiveDataSource.cpp b/velox/connectors/hive/HiveDataSource.cpp
@@ -513,6 +513,9 @@ std::unordered_map<std::string, RuntimeCounter> HiveDataSource::runtimeStats() {
   if (numBucketConversion_ > 0) {
     res.insert({"numBucketConversion", RuntimeCounter(numBucketConversion_)});
   }
+  for (const auto& storageStats : ioStats_->storageStats()) {
+    res.insert({storageStats.first, storageStats.second});
+  }
   return res;
 }
 

diff --git a/velox/connectors/hive/storage_adapters/abfs/AbfsFileSystem.h b/velox/connectors/hive/storage_adapters/abfs/AbfsFileSystem.h
@@ -40,7 +40,8 @@ class AbfsFileSystem : public FileSystem {
 
   std::unique_ptr<ReadFile> openFileForRead(
       std::string_view path,
-      const FileOptions& options = {}) override;
+      const FileOptions& options = {},
+      io::IoStatistics* ioStats = {}) override;
 
   std::unique_ptr<WriteFile> openFileForWrite(
       std::string_view path,

diff --git a/velox/connectors/hive/storage_adapters/gcs/GcsFileSystem.h b/velox/connectors/hive/storage_adapters/gcs/GcsFileSystem.h
@@ -41,7 +41,8 @@ class GcsFileSystem : public FileSystem {
   /// [[https://cloud.google.com/storage/docs/samples/storage-stream-file-download]].
   std::unique_ptr<ReadFile> openFileForRead(
       std::string_view path,
-      const FileOptions& options = {}) override;
+      const FileOptions& options = {},
+      io::IoStatistics* ioStats = {}) override;
 
   /// Initialize a WriteFile
   /// First the method google::cloud::storage::Client::GetObjectMetadata

diff --git a/velox/connectors/hive/storage_adapters/hdfs/HdfsFileSystem.h b/velox/connectors/hive/storage_adapters/hdfs/HdfsFileSystem.h
@@ -55,7 +55,8 @@ class HdfsFileSystem : public FileSystem {
 
   std::unique_ptr<ReadFile> openFileForRead(
       std::string_view path,
-      const FileOptions& options = {}) override;
+      const FileOptions& options = {},
+      io::IoStatistics* ioStats = {}) override;
 
   std::unique_ptr<WriteFile> openFileForWrite(
       std::string_view path,

diff --git a/velox/connectors/hive/storage_adapters/s3fs/S3FileSystem.h b/velox/connectors/hive/storage_adapters/s3fs/S3FileSystem.h
@@ -37,7 +37,8 @@ class S3FileSystem : public FileSystem {
 
   std::unique_ptr<ReadFile> openFileForRead(
       std::string_view s3Path,
-      const FileOptions& options = {}) override;
+      const FileOptions& options = {},
+      io::IoStatistics* ioStats = {}) override;
 
   std::unique_ptr<WriteFile> openFileForWrite(
       std::string_view s3Path,