feat: Implement idle source handling for WatermarkStatus

velox/connectors/kafka/KafkaDataSource.cpp

@@ -30,6 +30,7 @@ namespace facebook::velox::connector::kafka {
 namespace {
 constexpr const char* kTaskIndex = "task_index";
 constexpr const char* kTaskParallelism = "parallelism";
+constexpr uint32_t kPollTimeoutBackoffMillis = 1;
 
 uint32_t javaStringHashCode(const std::string& value) {
   uint32_t hash = 0;
@@ -58,6 +59,7 @@ KafkaDataSource::KafkaDataSource(
       config_(config),
       outputType_(outputType),
       batchSize_(config_->getDataBatchSize()) {
+  scheduler_.start();
   VELOX_CHECK(batchSize_ > 0, "Batch size config value must greater than 0.");
   const std::shared_ptr<KafkaTableHandle> kafkaTableHandle =
       std::dynamic_pointer_cast<KafkaTableHandle>(tableHandle);
@@ -80,6 +82,11 @@ KafkaDataSource::KafkaDataSource(
   createRecordDeserializer(config_->getFormat(), outputType_);
 }
 
+KafkaDataSource::~KafkaDataSource() {
+  scheduler_.shutdown();
+  completeBlockingFuture();
+}
+
 bool KafkaDataSource::consumerCanbeCreated() {
   return config_->exists(ConnectionConfig::kBootstrapServers) &&
       config_->exists(ConnectionConfig::kClientId) &&
@@ -230,9 +237,29 @@ void KafkaDataSource::applyTaskScopedClientId() {
         fmt::format("{}-{}", config_->getClientId(), taskIndex)}});
 }
 
+void KafkaDataSource::completeBlockingFuture() {
+  if (blockingPromise_.has_value()) {
+    blockingPromise_->setValue();
+    blockingPromise_.reset();
+  }
+}
+
+std::optional<RowVectorPtr> KafkaDataSource::blockOnPollTimeout(
+    velox::ContinueFuture& future) {
+  auto [promise, blockedFuture] =
+      makeVeloxContinuePromiseContract("KafkaDataSource::next");
+  blockingPromise_ = std::move(promise);
+  scheduler_.addFunctionOnce(
+      [this]() { completeBlockingFuture(); },
+      fmt::format("KafkaDataSource::next.{}", ++blockingSequence_),
+      std::chrono::milliseconds(kPollTimeoutBackoffMillis));
+  future = std::move(blockedFuture);
+  return std::nullopt;
+}
+
 std::optional<RowVectorPtr> KafkaDataSource::next(
     uint64_t,
-    velox::ContinueFuture&) {
+    velox::ContinueFuture& future) {
   std::optional<RowVectorPtr> res;
   size_t consumedMsgBytes = 0;
   if (queue_.empty()) {
@@ -242,7 +269,7 @@ std::optional<RowVectorPtr> KafkaDataSource::next(
     consumePos_ = 0;
     // If nothing consumed, return directly.
     if (consumedMsgBytes == 0) {
-      return res;
+      return blockOnPollTimeout(future);
     }
   }
   outRow_->prepareForReuse();

velox/connectors/kafka/KafkaDataSource.h

@@ -16,7 +16,9 @@
 #pragma once
 
 #include <cppkafka/cppkafka.h>
+#include <folly/experimental/FunctionScheduler.h>
 #include <map>
+#include <optional>
 #include "velox/common/base/RuntimeMetrics.h"
 #include "velox/common/future/VeloxPromise.h"
 #include "velox/connectors/Connector.h"
@@ -42,6 +44,8 @@ class KafkaDataSource : public DataSource {
       const ConnectorQueryCtx* connectorQueryCtx,
       const ConnectionConfigPtr& connectionConfig);
 
+  ~KafkaDataSource() override;
+
   /// Create a kafka connection to the given topics and partitions.
   void addSplit(ConnectorSplitPtr split) override;
 
@@ -98,6 +102,9 @@ class KafkaDataSource : public DataSource {
   KafkaConsumerPtr consumer_;
   /// The kafka record deserializer.
   KafkaRecordDeserializerPtr deserializer_;
+  folly::FunctionScheduler scheduler_;
+  std::optional<ContinuePromise> blockingPromise_;
+  uint64_t blockingSequence_{0};
   /// Count how many rows consumed.
   uint64_t completedRows_ = 0;
   /// Count how many bytes consumed.
@@ -152,6 +159,10 @@ class KafkaDataSource : public DataSource {
 
   void applyRestoredOffsets(cppkafka::TopicPartitionList& topicPartitions);
 
+  void completeBlockingFuture();
+
+  std::optional<RowVectorPtr> blockOnPollTimeout(velox::ContinueFuture& future);
+
   bool hasOffsetsToCommit(const TopicPartitionOffsets& offsets) const;
 
   void updateCommittedOffsets(const TopicPartitionOffsets& offsets);

velox/experimental/stateful/CMakeLists.txt

@@ -38,7 +38,8 @@ velox_add_library(
   WatermarkGenerator.cpp
   WatermarkSource.cpp
   WatermarkAssigner.cpp
-  StreamRecordTimestampInserter.cpp)
+  StreamRecordTimestampInserter.cpp
+  WatermarkIdleTracker.cpp)
 
 velox_link_libraries(
   velox_stateful_exec

velox/experimental/stateful/IdleTimerManager.h

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <atomic>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <utility>
+
+#include "velox/experimental/stateful/ProcessingTimeScheduler.h"
+
+namespace facebook::velox::stateful {
+
+// Shared one-shot idle timer helper for watermark-generating operators. It
+// owns the background scheduler, prevents duplicate timer registrations, skips
+// scheduling after the owner is already idle, and clamps past idle deadlines to
+// the current time.
+class IdleTimerManager {
+ public:
+  using TimerCallback = std::function<void(int64_t)>;
+
+  explicit IdleTimerManager(
+      std::unique_ptr<ProcessingTimeScheduler> scheduler = nullptr)
+      : scheduler_(std::move(scheduler)) {}
+
+  ~IdleTimerManager() {
+    shutdown();
+  }
+
+  void schedule(
+      int64_t now,
+      bool enabled,
+      bool idle,
+      int64_t idleDeadline,
+      TimerCallback callback) {
+    if (!enabled || idle || timerPending_.exchange(true)) {
+      return;
+    }
+    if (!scheduler_) {
+      scheduler_ = std::make_unique<SystemProcessingTimeScheduler>();
+    }
+    const int64_t fireAt = idleDeadline < now ? now : idleDeadline;
+    scheduler_->registerTimer(
+        fireAt,
+        ProcessingTimerTask(
+            fireAt,
+            [this, callback = std::move(callback)](int64_t timestamp) mutable {
+              timerPending_ = false;
+              callback(timestamp);
+            }));
+  }
+
+  void shutdown() {
+    if (scheduler_) {
+      scheduler_->close();
+      scheduler_.reset();
+    }
+    timerPending_ = false;
+  }
+
+ private:
+  std::unique_ptr<ProcessingTimeScheduler> scheduler_;
+  std::atomic<bool> timerPending_{false};
+};
+
+} // namespace facebook::velox::stateful

velox/experimental/stateful/StatefulOperator.cpp

@@ -210,6 +210,12 @@ void StatefulOperator::processWatermarkStatus(bool idle) {
   processWatermarkStatus(idle, 0);
 }
 
+void StatefulOperator::checkWatermarkStatus(int64_t now) {
+  for (auto& target : targets_) {
+    target->checkWatermarkStatus(now);
+  }
+}
+
 void StatefulOperator::initializeStateBackend(StateBackend* stateBackend) {
   if (!stateHandler_) {
     stateHandler_ = std::make_shared<StreamOperatorStateHandler>(

velox/experimental/stateful/StatefulOperator.h

@@ -50,6 +50,8 @@ class StatefulOperator {
     sink = operator_->operatorType() == "TableWrite";
   }
 
+  virtual ~StatefulOperator() = default;
+
   virtual void initialize();
 
   virtual bool isFinished();
@@ -74,6 +76,13 @@ class StatefulOperator {
 
   virtual void processWatermarkStatus(bool idle);
 
+  /// Called periodically by the driver thread (via StatefulTask::next) after
+  /// each drain attempt. Operators with idle-timeout detection (e.g.
+  /// WatermarkAssigner) override this to check for input idleness and emit
+  /// WatermarkStatus events into the task output. The default implementation
+  /// recursively forwards the call to downstream targets.
+  virtual void checkWatermarkStatus(int64_t now);
+
   virtual void initializeState();
 
   void initializeStateBackend(StateBackend* stateBackend);

velox/experimental/stateful/StatefulTask.cpp

@@ -22,6 +22,7 @@
 #include "velox/experimental/stateful/StatefulPlanner.h"
 #include "velox/experimental/stateful/state/HashMapStateBackend.h"
 #include "velox/experimental/stateful/state/RocksDBStateBackend.h"
+#include "velox/experimental/stateful/window/TimeWindowUtil.h"
 
 namespace facebook::velox::stateful {
 
@@ -141,21 +142,31 @@ StreamElementPtr StatefulTask::next(ContinueFuture* future, int32_t& retCode) {
 
   operatorChain_->advanceWithFuture(future);
   if (future != nullptr && future->valid()) {
-    return nullptr;
-  }
-  if (pendings_.empty()) {
-    if (operatorChain_->isFinished()) {
-      operatorChain_->finish();
-      finish();
-      // finish may trigger window flush and generate output.
+    // Source is blocked. Give operators a chance to detect idle input and
+    // emit WatermarkStatus events into pendings_ before returning.
+    checkWatermarkStatus(TimeWindowUtil::getCurrentProcessingTime());
+    if (pendings_.empty()) {
+      return nullptr;
+    }
+  } else {
+    if (pendings_.empty()) {
+      // No output produced. Give operators a chance to detect idle input.
+      checkWatermarkStatus(TimeWindowUtil::getCurrentProcessingTime());
       if (pendings_.empty()) {
-        retCode = 1;
-        return nullptr;
+        if (operatorChain_->isFinished()) {
+          operatorChain_->finish();
+          finish();
+          // finish may trigger window flush and generate output.
+          if (pendings_.empty()) {
+            retCode = 1;
+            return nullptr;
+          }
+        } else if (operatorChain_->sourceEmpty()) {
+          return nullptr;
+        } else {
+          return nullptr;
+        }
       }
-    } else if (operatorChain_->sourceEmpty()) {
-      return nullptr;
-    } else {
-      return nullptr;
     }
   }
   return popOutput();
@@ -201,6 +212,12 @@ void StatefulTask::notifyWatermarkStatus(bool idle) {
   operatorChain_->processWatermarkStatus(idle);
 }
 
+void StatefulTask::checkWatermarkStatus(int64_t now) {
+  if (operatorChain_) {
+    operatorChain_->checkWatermarkStatus(now);
+  }
+}
+
 void StatefulTask::initializeState(
     const std::shared_ptr<const KeyedStateBackendParameters> parameters,
     const std::vector<std::string>& checkpointRecords) {

velox/experimental/stateful/StatefulTask.h

@@ -72,6 +72,10 @@ class StatefulTask : public exec::Task {
 
   void notifyWatermarkStatus(bool idle);
 
+  /// Called on the driver thread to give each operator in the chain a chance
+  /// to perform idle-timeout detection and emit WatermarkStatus events.
+  void checkWatermarkStatus(int64_t now);
+
   void initializeState(
       const std::shared_ptr<const KeyedStateBackendParameters> params,
       const std::vector<std::string>& checkpointRecords = {});

velox/experimental/stateful/WatermarkAssigner.cpp

@@ -16,6 +16,7 @@
 #include "velox/experimental/stateful/WatermarkAssigner.h"
 
 #include "velox/experimental/stateful/WatermarkGenerator.h"
+#include "velox/experimental/stateful/window/TimeWindowUtil.h"
 
 namespace facebook::velox::stateful {
 
@@ -26,14 +27,26 @@ WatermarkAssigner::WatermarkAssigner(
     const int rowtimeFieldIndex,
     const int64_t watermarkInterval)
     : StatefulOperator(std::move(op), std::move(targets)),
-      idleTimeout_(idleTimeout),
       rowtimeFieldIndex_(rowtimeFieldIndex),
-      watermarkInterval_(watermarkInterval) {}
+      watermarkInterval_(watermarkInterval),
+      idleTracker_(idleTimeout) {}
+
+WatermarkAssigner::~WatermarkAssigner() {
+  shutdownScheduler();
+}
 
 void WatermarkAssigner::addInput(StreamElementPtr input) {
   auto record = std::static_pointer_cast<StreamRecord>(input);
   input_ = record->record();
   op()->addInput(input_);
+
+  if (idleTracker_.isEnabled()) {
+    int64_t now = TimeWindowUtil::getCurrentProcessingTime();
+    if (idleTracker_.onRecord(now)) {
+      // Was idle; emit WatermarkStatus.ACTIVE downstream.
+      emitWatermarkStatus(false);
+    }
+  }
 }
 
 void WatermarkAssigner::advance() {
@@ -78,7 +91,42 @@ void WatermarkAssigner::advance() {
   input_.reset();
 }
 
+void WatermarkAssigner::checkWatermarkStatus(int64_t now) {
+  if (!idleTracker_.isEnabled()) {
+    StatefulOperator::checkWatermarkStatus(now);
+    return;
+  }
+
+  if (idleTracker_.checkIdle(now)) {
+    emitWatermarkStatus(true);
+  }
+
+  scheduleIdleTimer(now);
+
+  // Forward to downstream targets.
+  StatefulOperator::checkWatermarkStatus(now);
+}
+
+void WatermarkAssigner::scheduleIdleTimer(int64_t now) {
+  idleTimerManager_.schedule(
+      now,
+      idleTracker_.isEnabled(),
+      idleTracker_.isIdle(),
+      idleTracker_.idleDeadline(),
+      [this](int64_t timestamp) {
+        auto bridge = nativeCallbackBridge();
+        if (bridge) {
+          bridge->onProcessingTime(timestamp);
+        }
+      });
+}
+
+void WatermarkAssigner::shutdownScheduler() {
+  idleTimerManager_.shutdown();
+}
+
 void WatermarkAssigner::close() {
+  shutdownScheduler();
   StatefulOperator::close();
   input_.reset();
 }