feat: Add storage statistics in IoStatistics

velox/common/io/IoStatistics.cpp

@@ -108,6 +108,24 @@ IoStatistics::operationStats() const {
   return operationStats_;
 }
 
+std::unordered_map<std::string, RuntimeMetric> IoStatistics::storageStats()
+    const {
+  std::lock_guard<std::mutex> lock{storageStatsMutex_};
+  return storageStats_;
+}
+
+void IoStatistics::addStorageStats(
+    const std::string& name,
+    const RuntimeCounter& counter) {
+  std::lock_guard<std::mutex> lock{storageStatsMutex_};
+  if (storageStats_.count(name) == 0) {
+    storageStats_.emplace(name, RuntimeMetric(counter.unit));
+  } else {
+    VELOX_CHECK_EQ(storageStats_.at(name).unit, counter.unit);
+  }
+  storageStats_.at(name).addValue(counter.value);
+}
+
 void IoStatistics::merge(const IoStatistics& other) {
   rawBytesRead_ += other.rawBytesRead_;
   rawBytesWritten_ += other.rawBytesWritten_;
@@ -119,9 +137,20 @@ void IoStatistics::merge(const IoStatistics& other) {
   ramHit_.merge(other.ramHit_);
   ssdRead_.merge(other.ssdRead_);
   queryThreadIoLatency_.merge(other.queryThreadIoLatency_);
-  std::lock_guard<std::mutex> l(operationStatsMutex_);
-  for (auto& item : other.operationStats_) {
-    operationStats_[item.first].merge(item.second);
+  {
+    const auto& otherOperationStats = other.operationStats();
+    std::lock_guard<std::mutex> l(operationStatsMutex_);
+    for (auto& item : otherOperationStats) {
+      operationStats_[item.first].merge(item.second);
+    }
+  }
+
+  {
+    const auto& otherStorageStats = other.storageStats();
+    std::lock_guard<std::mutex> storageStatsLock(storageStatsMutex_);
+    for (auto& item : otherStorageStats) {
+      storageStats_[item.first].merge(item.second);
+    }
   }
 }
 

velox/common/io/IoStatistics.h

@@ -23,6 +23,8 @@
 #include <unordered_map>
 
 #include <folly/dynamic.h>
+#include "velox/common/base/Exceptions.h"
+#include "velox/common/base/RuntimeMetrics.h"
 
 namespace facebook::velox::io {
 
@@ -140,6 +142,9 @@ class IoStatistics {
       const uint64_t partialThrottleCount = 0);
 
   std::unordered_map<std::string, OperationCounters> operationStats() const;
+  std::unordered_map<std::string, RuntimeMetric> storageStats() const;
+
+  void addStorageStats(const std::string& name, const RuntimeCounter& counter);
 
   void merge(const IoStatistics& other);
 
@@ -172,7 +177,9 @@ class IoStatistics {
   IoCounter queryThreadIoLatency_;
 
   std::unordered_map<std::string, OperationCounters> operationStats_;
+  std::unordered_map<std::string, RuntimeMetric> storageStats_;
   mutable std::mutex operationStatsMutex_;
+  mutable std::mutex storageStatsMutex_;
 };
 
 } // namespace facebook::velox::io