finish io

fix

Author: FelixYBW

velox/common/io/IoStatistics.h

@@ -127,6 +127,34 @@ class IoStatistics {
     return queryThreadIoLatency_;
   }
 
+  void startIO() {
+    std::lock_guard<std::mutex> lock(mtx_);
+    if (active_requests_ == 0 && tracking_idle_) {
+      auto now = std::chrono::steady_clock::now();
+      queryThreadIoLatency_.increment(std::chrono::duration_cast<std::chrono::milliseconds>(now - idle_start_).count()*1000);
+      tracking_idle_ = false;
+    }
+    ++active_requests_;
+  }
+  
+  void endIO() {
+    std::lock_guard<std::mutex> lock(mtx_);
+    --active_requests_;
+    if (active_requests_ == 0) {
+      idle_start_ = std::chrono::steady_clock::now();
+      tracking_idle_ = true;
+    }
+  }
+
+  void finish() {
+    std::lock_guard<std::mutex> lock(mtx_);
+    if (tracking_idle_) {
+      auto now = std::chrono::steady_clock::now();
+      queryThreadIoLatency_.increment(std::chrono::duration_cast<std::chrono::milliseconds>(now - idle_start_).count()*1000);
+      idle_start_ = now;
+    }
+  }  
+
   void incOperationCounters(
       const std::string& operation,
       const uint64_t resourceThrottleCount,
@@ -173,6 +201,13 @@ class IoStatistics {
 
   std::unordered_map<std::string, OperationCounters> operationStats_;
   mutable std::mutex operationStatsMutex_;
+
+  std::mutex mtx_;
+  int active_requests_{0};
+  std::chrono::steady_clock::time_point idle_start_;
+  long idle_time_{0};
+  bool tracking_idle_{false};  
+
 };
 
 } // namespace facebook::velox::io

velox/connectors/hive/HiveDataSource.cpp

@@ -421,6 +421,7 @@ void HiveDataSource::addDynamicFilter(
 std::unordered_map<std::string, RuntimeMetric>
 HiveDataSource::getRuntimeStats() {
   auto res = runtimeStats_.toRuntimeMetricMap();
+  ioStats_->finish();
   res.insert(
       {{"numPrefetch", RuntimeMetric(ioStats_->prefetch().count())},
        {"prefetchBytes",

velox/dwio/common/DirectBufferedInput.cpp

@@ -318,13 +318,14 @@ std::vector<cache::CachePin> DirectCoalescedLoad::loadData(bool prefetch) {
   uint64_t usecs = 0;
   {
     MicrosecondTimer timer(&usecs);
+    ioStats_->startIO();
     input_->read(buffers, requests_[0].region.offset, LogType::FILE);
+    ioStats_->endIO();
   }
 
   ioStats_->read().increment(size + overread);
   ioStats_->incRawBytesRead(size);
   ioStats_->incTotalScanTime(usecs * 1'000);
-  ioStats_->queryThreadIoLatency().increment(usecs);
   ioStats_->incRawOverreadBytes(overread);
   if (prefetch) {
     ioStats_->prefetch().increment(size + overread);

velox/dwio/common/DirectInputStream.cpp

@@ -150,10 +150,11 @@ void DirectInputStream::loadSync() {
   uint64_t usecs = 0;
   {
     MicrosecondTimer timer(&usecs);
+    ioStats_->startIO();
     input_->read(ranges, loadedRegion_.offset, LogType::FILE);
+    ioStats_->endIO();
   }
   ioStats_->read().increment(loadedRegion_.length);
-  ioStats_->queryThreadIoLatency().increment(usecs);
   ioStats_->incTotalScanTime(usecs * 1'000);
 }
 
@@ -173,7 +174,6 @@ void DirectInputStream::loadPosition() {
         loadedRegion_.offset = region_.offset;
         loadedRegion_.length = load->getData(region_.offset, data_, tinyData_);
       }
-      ioStats_->queryThreadIoLatency().increment(loadUs);
     } else {
       // Standalone stream, not part of coalesced load.
       loadedRegion_.offset = 0;