feat: Collect storage stat for Tables scan (facebookincubator#12210)

Summary: Pull Request resolved: facebookincubator#12210 Differential Revision: D68872215 Pulled By: kewang1024
kewang1024 · Jan 30, 2025 · d50959c · d50959c
1 parent 5f8e7f1
commit d50959c
Show file tree

Hide file tree

Showing 17 changed files with 245 additions and 96 deletions.
diff --git a/velox/common/file/File.cpp b/velox/common/file/File.cpp
@@ -57,17 +57,21 @@ T getAttribute(
 }
 } // namespace
 
-std::string ReadFile::pread(uint64_t offset, uint64_t length) const {
+std::string ReadFile::pread(
+    uint64_t offset,
+    uint64_t length,
+    io::IoStatistics* stats) const {
   std::string buf;
   buf.resize(length);
-  auto res = pread(offset, length, buf.data());
+  auto res = pread(offset, length, buf.data(), stats);
   buf.resize(res.size());
   return buf;
 }
 
 uint64_t ReadFile::preadv(
     uint64_t offset,
-    const std::vector<folly::Range<char*>>& buffers) const {
+    const std::vector<folly::Range<char*>>& buffers,
+    io::IoStatistics* stats) const {
   auto fileSize = size();
   uint64_t numRead = 0;
   if (offset >= fileSize) {
@@ -77,7 +81,7 @@ uint64_t ReadFile::preadv(
     auto copySize = std::min<size_t>(range.size(), fileSize - offset);
     // NOTE: skip the gap in case of coalesce io.
     if (range.data() != nullptr) {
-      pread(offset, copySize, range.data());
+      pread(offset, copySize, range.data(), stats);
     }
     offset += copySize;
     numRead += copySize;
@@ -87,29 +91,36 @@ uint64_t ReadFile::preadv(
 
 uint64_t ReadFile::preadv(
     folly::Range<const common::Region*> regions,
-    folly::Range<folly::IOBuf*> iobufs) const {
+    folly::Range<folly::IOBuf*> iobufs,
+    io::IoStatistics* stats) const {
   VELOX_CHECK_EQ(regions.size(), iobufs.size());
   uint64_t length = 0;
   for (size_t i = 0; i < regions.size(); ++i) {
     const auto& region = regions[i];
     auto& output = iobufs[i];
     output = folly::IOBuf(folly::IOBuf::CREATE, region.length);
-    pread(region.offset, region.length, output.writableData());
+    pread(region.offset, region.length, output.writableData(), stats);
     output.append(region.length);
     length += region.length;
   }
 
   return length;
 }
 
-std::string_view
-InMemoryReadFile::pread(uint64_t offset, uint64_t length, void* buf) const {
+std::string_view InMemoryReadFile::pread(
+    uint64_t offset,
+    uint64_t length,
+    void* buf,
+    io::IoStatistics* stats) const {
   bytesRead_ += length;
   memcpy(buf, file_.data() + offset, length);
   return {static_cast<char*>(buf), length};
 }
 
-std::string InMemoryReadFile::pread(uint64_t offset, uint64_t length) const {
+std::string InMemoryReadFile::pread(
+    uint64_t offset,
+    uint64_t length,
+    io::IoStatistics* stats) const {
   bytesRead_ += length;
   return std::string(file_.data() + offset, length);
 }
@@ -187,15 +198,19 @@ void LocalReadFile::preadInternal(uint64_t offset, uint64_t length, char* pos)
       folly::errnoStr(errno));
 }
 
-std::string_view
-LocalReadFile::pread(uint64_t offset, uint64_t length, void* buf) const {
+std::string_view LocalReadFile::pread(
+    uint64_t offset,
+    uint64_t length,
+    void* buf,
+    io::IoStatistics* stats) const {
   preadInternal(offset, length, static_cast<char*>(buf));
   return {static_cast<char*>(buf), length};
 }
 
 uint64_t LocalReadFile::preadv(
     uint64_t offset,
-    const std::vector<folly::Range<char*>>& buffers) const {
+    const std::vector<folly::Range<char*>>& buffers,
+    io::IoStatistics* stats) const {
   // Dropped bytes sized so that a typical dropped range of 50K is not
   // too many iovecs.
   static thread_local std::vector<char> droppedBytes(16 * 1024);
@@ -251,16 +266,18 @@ uint64_t LocalReadFile::preadv(
 
 folly::SemiFuture<uint64_t> LocalReadFile::preadvAsync(
     uint64_t offset,
-    const std::vector<folly::Range<char*>>& buffers) const {
+    const std::vector<folly::Range<char*>>& buffers,
+    io::IoStatistics* stats) const {
   if (!executor_) {
-    return ReadFile::preadvAsync(offset, buffers);
+    return ReadFile::preadvAsync(offset, buffers, stats);
   }
   auto [promise, future] = folly::makePromiseContract<uint64_t>();
   executor_->add([this,
                   _promise = std::move(promise),
                   _offset = offset,
-                  _buffers = buffers]() mutable {
-    auto delegateFuture = ReadFile::preadvAsync(_offset, _buffers);
+                  _buffers = buffers,
+                  _stats = stats]() mutable {
+    auto delegateFuture = ReadFile::preadvAsync(_offset, _buffers, _stats);
     _promise.setTry(std::move(delegateFuture).getTry());
   });
   return std::move(future);

diff --git a/velox/common/file/File.h b/velox/common/file/File.h
@@ -40,6 +40,7 @@
 
 #include "velox/common/base/Exceptions.h"
 #include "velox/common/file/Region.h"
+#include "velox/common/io/IoStatistics.h"
 
 namespace facebook::velox {
 
@@ -52,13 +53,19 @@ class ReadFile {
   // buffer 'buf'. The bytes are returned as a string_view pointing to 'buf'.
   //
   // This method should be thread safe.
-  virtual std::string_view pread(uint64_t offset, uint64_t length, void* buf)
-      const = 0;
+  virtual std::string_view pread(
+      uint64_t offset,
+      uint64_t length,
+      void* buf,
+      io::IoStatistics* stats = nullptr) const = 0;
 
   // Same as above, but returns owned data directly.
   //
   // This method should be thread safe.
-  virtual std::string pread(uint64_t offset, uint64_t length) const;
+  virtual std::string pread(
+      uint64_t offset,
+      uint64_t length,
+      io::IoStatistics* stats = nullptr) const;
 
   // Reads starting at 'offset' into the memory referenced by the
   // Ranges in 'buffers'. The buffers are filled left to right. A
@@ -67,7 +74,8 @@ class ReadFile {
   // This method should be thread safe.
   virtual uint64_t preadv(
       uint64_t /*offset*/,
-      const std::vector<folly::Range<char*>>& /*buffers*/) const;
+      const std::vector<folly::Range<char*>>& /*buffers*/,
+      io::IoStatistics* stats = nullptr) const;
 
   // Vectorized read API. Implementations can coalesce and parallelize.
   // The offsets don't need to be sorted.
@@ -82,7 +90,8 @@ class ReadFile {
   // This method should be thread safe.
   virtual uint64_t preadv(
       folly::Range<const common::Region*> regions,
-      folly::Range<folly::IOBuf*> iobufs) const;
+      folly::Range<folly::IOBuf*> iobufs,
+      io::IoStatistics* stats = nullptr) const;
 
   /// Like preadv but may execute asynchronously and returns the read size or
   /// exception via SemiFuture. Use hasPreadvAsync() to check if the
@@ -91,9 +100,10 @@ class ReadFile {
   /// This method should be thread safe.
   virtual folly::SemiFuture<uint64_t> preadvAsync(
       uint64_t offset,
-      const std::vector<folly::Range<char*>>& buffers) const {
+      const std::vector<folly::Range<char*>>& buffers,
+      io::IoStatistics* stats = nullptr) const {
     try {
-      return folly::SemiFuture<uint64_t>(preadv(offset, buffers));
+      return folly::SemiFuture<uint64_t>(preadv(offset, buffers, stats));
     } catch (const std::exception& e) {
       return folly::makeSemiFuture<uint64_t>(e);
     }
@@ -213,10 +223,14 @@ class InMemoryReadFile : public ReadFile {
   explicit InMemoryReadFile(std::string file)
       : ownedFile_(std::move(file)), file_(ownedFile_) {}
 
-  std::string_view pread(uint64_t offset, uint64_t length, void* buf)
-      const override;
+  std::string_view pread(
+      uint64_t offset,
+      uint64_t length,
+      void* buf,
+      io::IoStatistics* stats = nullptr) const override;
 
-  std::string pread(uint64_t offset, uint64_t length) const override;
+  std::string pread(uint64_t offset, uint64_t length, io::IoStatistics* stats)
+      const override;
 
   uint64_t size() const final {
     return file_.size();
@@ -278,18 +292,23 @@ class LocalReadFile final : public ReadFile {
 
   ~LocalReadFile();
 
-  std::string_view pread(uint64_t offset, uint64_t length, void* buf)
-      const final;
+  std::string_view pread(
+      uint64_t offset,
+      uint64_t length,
+      void* buf,
+      io::IoStatistics* stats = nullptr) const final;
 
   uint64_t size() const final;
 
   uint64_t preadv(
       uint64_t offset,
-      const std::vector<folly::Range<char*>>& buffers) const final;
+      const std::vector<folly::Range<char*>>& buffers,
+      io::IoStatistics* stats = nullptr) const final;
 
   folly::SemiFuture<uint64_t> preadvAsync(
       uint64_t offset,
-      const std::vector<folly::Range<char*>>& buffers) const override;
+      const std::vector<folly::Range<char*>>& buffers,
+      io::IoStatistics* stats = nullptr) const override;
 
   bool hasPreadvAsync() const override {
     return executor_ != nullptr;

diff --git a/velox/common/file/tests/FaultyFile.cpp b/velox/common/file/tests/FaultyFile.cpp
@@ -30,45 +30,52 @@ FaultyReadFile::FaultyReadFile(
   VELOX_CHECK_NOT_NULL(delegatedFile_);
 }
 
-std::string_view
-FaultyReadFile::pread(uint64_t offset, uint64_t length, void* buf) const {
+std::string_view FaultyReadFile::pread(
+    uint64_t offset,
+    uint64_t length,
+    void* buf,
+    io::IoStatistics* stats) const {
   if (injectionHook_ != nullptr) {
     FaultFileReadOperation op(path_, offset, length, buf);
     injectionHook_(&op);
     if (!op.delegate) {
       return std::string_view(static_cast<char*>(op.buf), op.length);
     }
   }
-  return delegatedFile_->pread(offset, length, buf);
+  return delegatedFile_->pread(offset, length, buf, stats);
 }
 
 uint64_t FaultyReadFile::preadv(
     uint64_t offset,
-    const std::vector<folly::Range<char*>>& buffers) const {
+    const std::vector<folly::Range<char*>>& buffers,
+    io::IoStatistics* stats) const {
   if (injectionHook_ != nullptr) {
     FaultFileReadvOperation op(path_, offset, buffers);
     injectionHook_(&op);
     if (!op.delegate) {
       return op.readBytes;
     }
   }
-  return delegatedFile_->preadv(offset, buffers);
+  return delegatedFile_->preadv(offset, buffers, stats);
 }
 
 folly::SemiFuture<uint64_t> FaultyReadFile::preadvAsync(
     uint64_t offset,
-    const std::vector<folly::Range<char*>>& buffers) const {
+    const std::vector<folly::Range<char*>>& buffers,
+    io::IoStatistics* stats) const {
   // TODO: add fault injection for async read later.
   if (delegatedFile_->hasPreadvAsync() || executor_ == nullptr) {
-    return delegatedFile_->preadvAsync(offset, buffers);
+    return delegatedFile_->preadvAsync(offset, buffers, stats);
   }
   auto promise = std::make_unique<folly::Promise<uint64_t>>();
   folly::SemiFuture<uint64_t> future = promise->getSemiFuture();
   executor_->add([this,
                   _promise = std::move(promise),
                   _offset = offset,
-                  _buffers = buffers]() {
-    auto delegateFuture = delegatedFile_->preadvAsync(_offset, _buffers);
+                  _buffers = buffers,
+                  _stats = stats]() {
+    auto delegateFuture =
+        delegatedFile_->preadvAsync(_offset, _buffers, _stats);
     _promise->setValue(delegateFuture.wait().value());
   });
   return future;

diff --git a/velox/common/file/tests/FaultyFile.h b/velox/common/file/tests/FaultyFile.h
@@ -35,12 +35,16 @@ class FaultyReadFile : public ReadFile {
     return delegatedFile_->size();
   }
 
-  std::string_view pread(uint64_t offset, uint64_t length, void* buf)
-      const override;
+  std::string_view pread(
+      uint64_t offset,
+      uint64_t length,
+      void* buf,
+      io::IoStatistics* stats = nullptr) const override;
 
   uint64_t preadv(
       uint64_t offset,
-      const std::vector<folly::Range<char*>>& buffers) const override;
+      const std::vector<folly::Range<char*>>& buffers,
+      io::IoStatistics* stats = nullptr) const override;
 
   uint64_t memoryUsage() const override {
     return delegatedFile_->memoryUsage();
@@ -67,7 +71,8 @@ class FaultyReadFile : public ReadFile {
 
   folly::SemiFuture<uint64_t> preadvAsync(
       uint64_t offset,
-      const std::vector<folly::Range<char*>>& buffers) const override;
+      const std::vector<folly::Range<char*>>& buffers,
+      io::IoStatistics* stats = nullptr) const override;
 
  private:
   const std::string path_;

diff --git a/velox/common/io/IoStatistics.cpp b/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);
+    }
   }
 }