async ca: rework watermark handling and fix inputtransform and watermark interaction

ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp

@@ -181,9 +181,7 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
         DUMP(ProcessOutputsState, LastPopReturnedNoData);
 
         html << "<h3>Watermarks</h3>";
-        for (const auto& [time, id]: WatermarkTakeInputChannelDataRequests) {
-            html << "WatermarkTakeInputChannelDataRequests: " << time.ToString() << " " << id << "<br />";
-        }
+        DUMP(WatermarksTracker, GetPendingWatermark, ());
 
         html << "<h3>CPU Quota</h3>";
         html << "QuoterServiceActorId: " << QuoterServiceActorId.ToString() << "<br />";
@@ -250,7 +248,6 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
             DUMP(info, ChannelId);
             DUMP(info, SrcStageId);
             DUMP(info, HasPeer);
-            html << "PendingWatermarks: " << !info.PendingWatermarks.empty() << " " << (info.PendingWatermarks.empty() ? TString{} : info.PendingWatermarks.back().ToString()) << "<br />";
             html << "WatermarksMode: " << NDqProto::EWatermarksMode_Name(info.WatermarksMode) << "<br />";
             html << "PendingCheckpoint: " << info.PendingCheckpoint.has_value() << " " << (info.PendingCheckpoint ? TStringBuilder{} << info.PendingCheckpoint->GetId() << " " << info.PendingCheckpoint->GetGeneration() : TString{}) << "<br />";
             html << "CheckpointingMode: " << NDqProto::ECheckpointingMode_Name(info.CheckpointingMode) << "<br />";
@@ -309,7 +306,6 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
                 html << "DqInputBuffer.InputType: " << (buffer->GetInputType() ? buffer->GetInputType()->GetKindAsStr() : TString{"unknown"})  << "<br />";
                 html << "DqInputBuffer.InputWidth: " << (buffer->GetInputWidth() ? ToString(*buffer->GetInputWidth()) : TString{"unknown"})  << "<br />";
                 html << "DqInputBuffer.IsFinished: " << buffer->IsFinished() << "<br />";
-                html << "DqInputBuffer.IsPaused: " << buffer->IsPaused() << "<br />";
                 html << "DqInputBuffer.IsPending: " << buffer->IsPending() << "<br />";
 
                 const auto& popStats = buffer->GetPopStats();
@@ -524,7 +520,7 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
         }
 
         const bool wasFinished = outputChannel.Finished;
-        auto channelId = outputChannel.ChannelId;
+        const auto channelId = outputChannel.ChannelId;
 
         const auto& peerState = Channels->GetOutputChannelInFlightState(channelId);
 
@@ -619,18 +615,6 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
         TMaybe<TInstant> watermark;
         if (channelData.HasWatermark()) {
             watermark = TInstant::MicroSeconds(channelData.GetWatermark().GetTimestampUs());
-
-            const bool channelWatermarkChanged = WatermarksTracker.NotifyInChannelWatermarkReceived(
-                inputChannel->ChannelId,
-                *watermark
-            );
-
-            if (channelWatermarkChanged) {
-                CA_LOG_T("Pause input channel " << channelData.GetChannelId() << " bacause of watermark " << *watermark);
-                inputChannel->Pause(*watermark);
-            }
-
-            WatermarkTakeInputChannelDataRequests[*watermark]++;
         }
 
         TDqSerializedBatch batch;
@@ -646,7 +630,8 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
             channelData.GetChannelId(),
             batch.RowCount() ? std::optional{std::move(batch)} : std::nullopt,
             finished,
-            channelData.HasCheckpoint()
+            channelData.HasCheckpoint(),
+            watermark
         );
 
         Send(TaskRunnerActorId, ev.Release(), 0, Cookie);
@@ -680,22 +665,6 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
         TBase::PassAway();
     }
 
-    TMaybe<TInstant> GetWatermarkRequest() {
-        if (!WatermarksTracker.HasPendingWatermark()) {
-            return Nothing();
-        }
-
-        const auto pendingWatermark = *WatermarksTracker.GetPendingWatermark();
-        if (WatermarkTakeInputChannelDataRequests.contains(pendingWatermark)) {
-            // Not all precending to watermark input channels data has been injected
-            return Nothing();
-        }
-
-        MetricsReporter.ReportInjectedToTaskRunnerWatermark(pendingWatermark);
-
-        return pendingWatermark;
-    }
-
     TMaybe<NDqProto::TCheckpoint> GetCheckpointRequest() {
         if (!CheckpointRequestedFromTaskRunner && Checkpoints && Checkpoints->HasPendingCheckpoint() && !Checkpoints->ComputeActorStateSaved()) {
             CheckpointRequestedFromTaskRunner = true;
@@ -801,9 +770,22 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
             }
         }
 
-        if (ev->Get()->WatermarkInjectedToOutputs && !WatermarksTracker.HasOutputChannels()) {
-            ResumeInputsByWatermark(*WatermarksTracker.GetPendingWatermark());
-            WatermarksTracker.PopPendingWatermark();
+        if (auto watermark = ev->Get()->WatermarkInjectedToOutputs) {
+            ResumeInputsByWatermark(*watermark);
+            if (WatermarksTracker.GetPendingWatermark() <= watermark) {
+                MetricsReporter.ReportInjectedToOutputsWatermark(*watermark);
+                WatermarksTracker.PopPendingWatermark();
+            }
+        }
+
+        for (auto inputChannelId : ev->Get()->FinishedInputsWithWatermarks) {
+            CA_LOG_T("Unregister watermarked input channel " << inputChannelId);
+            WatermarksTracker.UnregisterInputChannel(inputChannelId);
+        }
+
+        for (auto sourceId : ev->Get()->FinishedSourcesWithWatermarks) {
+            CA_LOG_T("Unregister watermarked async input " << sourceId);
+            WatermarksTracker.UnregisterInputChannel(sourceId);
         }
 
         ReadyToCheckpointFlag = (bool) ev->Get()->ProgramState;
@@ -866,6 +848,7 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
         if (Stat) {
             Stat->AddCounters2(ev->Get()->Sensors);
         }
+        CA_LOG_T("OnOutputChannelData");
         auto it = OutputChannelsMap.find(ev->Get()->ChannelId);
         Y_ABORT_UNLESS(it != OutputChannelsMap.end());
         TOutputChannelInfo& outputChannel = it->second;
@@ -906,6 +889,14 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
     }
 
     void SendAsyncChannelData(TOutputChannelInfo& outputChannel) {
+        CA_LOG_T("SendAsyncChannelData. Channel: " << outputChannel.ChannelId
+            << " Finished: " << outputChannel.Finished
+            << " Data: { DataSize: " << outputChannel.AsyncData->Data.size()
+            << ", Watermark: " << outputChannel.AsyncData->Watermark
+            << ", Checkpoint: " << outputChannel.AsyncData->Checkpoint
+            << ", Finished: " << outputChannel.AsyncData->Finished
+            << ", Changed: " << outputChannel.AsyncData->Changed << "}"
+            );
         Y_ABORT_UNLESS(outputChannel.AsyncData);
 
         // If the channel has finished, then the data received after drain is no longer needed
@@ -922,18 +913,6 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
         ProcessOutputsState.HasDataToSend |= !outputChannel.Finished;
         ProcessOutputsState.LastPopReturnedNoData = asyncData.Data.empty();
 
-        if (asyncData.Watermark.Defined()) {
-            const auto watermark = TInstant::MicroSeconds(asyncData.Watermark->GetTimestampUs());
-            const bool shouldResumeInputs = WatermarksTracker.NotifyOutputChannelWatermarkSent(
-                outputChannel.ChannelId,
-                watermark
-            );
-
-            if (shouldResumeInputs) {
-                ResumeInputsByWatermark(watermark);
-            }
-        }
-
         if (!shouldSkipData) {
             if (asyncData.Checkpoint.Defined()) {
                 ResumeInputsByCheckpoint();
@@ -980,28 +959,32 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
     }
 
     void OnInputChannelDataAck(NTaskRunnerActor::TEvInputChannelDataAck::TPtr& ev) {
-        auto it = TakeInputChannelDataRequests.find(ev->Cookie);
+        const auto it = TakeInputChannelDataRequests.find(ev->Cookie);
         YQL_ENSURE(it != TakeInputChannelDataRequests.end());
 
+        const auto channelId = it->second.ChannelId;
+        const auto watermark = it->second.Watermark;
+
         CA_LOG_T("Input data push finished. Cookie: " << ev->Cookie
-            << " Watermark: " << it->second.Watermark
+            << " Watermark: " << watermark
             << " Ack: " << it->second.Ack
             << " TakeInputChannelDataRequests: " << TakeInputChannelDataRequests.size()
-            << " WatermarkTakeInputChannelDataRequests: " << WatermarkTakeInputChannelDataRequests.size());
-
-        if (it->second.Watermark.Defined()) {
-            auto& ct = WatermarkTakeInputChannelDataRequests.at(*it->second.Watermark);
-            if (--ct == 0) {
-                WatermarkTakeInputChannelDataRequests.erase(*it->second.Watermark);
-            }
-        }
+            );
 
-        TInputChannelInfo* inputChannel = InputChannelsMap.FindPtr(it->second.ChannelId);
+        TInputChannelInfo* inputChannel = InputChannelsMap.FindPtr(channelId);
         Y_ABORT_UNLESS(inputChannel);
+
         inputChannel->FreeSpace = ev->Get()->FreeSpace;
 
+        if (watermark) {
+            if (WatermarksTracker.NotifyInChannelWatermarkReceived( channelId, *watermark)) {
+                CA_LOG_T("Pause input channel " << channelId << " because of watermark");
+                inputChannel->Pause(*watermark); // XXX does nothing in async CA
+            }
+        }
+
         if (it->second.Ack) {
-            Channels->SendChannelDataAck(it->second.ChannelId, inputChannel->FreeSpace);
+            Channels->SendChannelDataAck(channelId, inputChannel->FreeSpace);
         }
 
         TakeInputChannelDataRequests.erase(it);
@@ -1141,14 +1124,6 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
             ContinueRunEvent = std::make_unique<NTaskRunnerActor::TEvContinueRun>();
         }
         ContinueRunEvent->CheckpointOnly = checkpointOnly;
-        if (TMaybe<TInstant> watermarkRequest = GetWatermarkRequest()) {
-            if (!ContinueRunEvent->WatermarkRequest) {
-                ContinueRunEvent->WatermarkRequest.ConstructInPlace();
-                ContinueRunEvent->WatermarkRequest->Watermark = *watermarkRequest;
-            } else {
-                ContinueRunEvent->WatermarkRequest->Watermark = Max(ContinueRunEvent->WatermarkRequest->Watermark, *watermarkRequest);
-            }
-        }
         if (checkpointRequest) {
             if (!ContinueRunEvent->CheckpointRequest) {
                 ContinueRunEvent->CheckpointRequest.ConstructInPlace(*checkpointRequest);
@@ -1157,6 +1132,11 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
                 Y_ABORT_UNLESS(ContinueRunEvent->CheckpointRequest->Checkpoint.GetId() == checkpointRequest->GetId());
             }
         }
+        if (auto watermarkRequest = WatermarksTracker.GetPendingWatermark()) {
+            Y_ENSURE(*watermarkRequest >= ContinueRunEvent->WatermarkRequest);
+            ContinueRunEvent->WatermarkRequest = *watermarkRequest;
+            MetricsReporter.ReportInjectedToTaskRunnerWatermark(*watermarkRequest);
+        }
 
         if (!UseCpuQuota()) {
             Send(TaskRunnerActorId, ContinueRunEvent.release());
@@ -1225,9 +1205,6 @@ class TDqAsyncComputeActor : public TDqComputeActorBase<TDqAsyncComputeActor, TC
         TMaybe<TInstant> Watermark;
     };
     THashMap<ui64, TTakeInputChannelData> TakeInputChannelDataRequests;
-    // Watermark should be injected to task runner only after all precending data is injected
-    // This hash map will help to track the right moment
-    THashMap<TInstant, ui32> WatermarkTakeInputChannelDataRequests;
     ui64 Cookie = 0;
     NDq::TDqTaskRunnerStatsView TaskRunnerStats;
     bool ReadyToCheckpointFlag;

ydb/library/yql/dq/actors/compute/dq_compute_actor_async_input_helper.h

@@ -45,8 +45,9 @@ struct TComputeActorAsyncInputHelper {
     }
 
     void ResumeByWatermark(TInstant watermark) {
-        YQL_ENSURE(watermark == PendingWatermark);
-        PendingWatermark = Nothing();
+        if (watermark >= PendingWatermark) {
+            PendingWatermark = Nothing();
+        }
     }
 
     virtual i64 GetFreeSpace() const = 0;

ydb/library/yql/dq/actors/compute/dq_compute_actor_async_io.h

@@ -55,6 +55,7 @@ enum class EResumeSource : ui32 {
     CADataSent,
     CAPendingOutput,
     CATaskRunnerCreated,
+    CAWatermarkInject,
 
     Last,
 };

ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h

@@ -764,6 +764,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
             sourceInfo.ResumeByWatermark(watermark);
         }
 
+        // XXX Does nothing in async CA, not used (yet) in sync CA
         for (auto& [id, channelInfo] : InputChannelsMap) {
             if (channelInfo.WatermarksMode == NDqProto::EWatermarksMode::WATERMARKS_MODE_DISABLED) {
                 continue;
@@ -839,7 +840,6 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
         ui32 SrcStageId;
         IDqInputChannel::TPtr Channel;
         bool HasPeer = false;
-        std::queue<TInstant> PendingWatermarks;
         const NDqProto::EWatermarksMode WatermarksMode;
         std::optional<NDqProto::TCheckpoint> PendingCheckpoint;
         const NDqProto::ECheckpointingMode CheckpointingMode;
@@ -860,38 +860,36 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
         }
 
         bool IsPaused() const {
-            return !PendingWatermarks.empty() || PendingCheckpoint.has_value();
+            return PendingCheckpoint.has_value();
         }
 
         void Pause(TInstant watermark) {
             YQL_ENSURE(WatermarksMode != NDqProto::WATERMARKS_MODE_DISABLED);
 
-            PendingWatermarks.emplace(watermark);
+            if (Channel) {
+                Channel->AddWatermark(watermark);
+            }
         }
 
         void Pause(const NDqProto::TCheckpoint& checkpoint) {
             YQL_ENSURE(!PendingCheckpoint);
             YQL_ENSURE(CheckpointingMode != NDqProto::CHECKPOINTING_MODE_DISABLED);
             PendingCheckpoint = checkpoint;
             if (Channel) {  // async actor doesn't hold channels, so channel is paused in task runner actor
-                Channel->Pause();
+                Channel->PauseByCheckpoint();
             }
         }
 
         void ResumeByWatermark(TInstant watermark) {
-            while (!PendingWatermarks.empty() && PendingWatermarks.front() <= watermark) {
-                if (PendingWatermarks.front() != watermark) {
-                    CA_LOG_W("Input channel " << ChannelId <<
-                        " watermarks were collapsed. See YQ-1441. Dropped watermark: " << PendingWatermarks.front());
-                }
-                PendingWatermarks.pop();
+            if (Channel) {
+                Channel->ResumeByWatermark(watermark);
             }
         }
 
         void ResumeByCheckpoint() {
             PendingCheckpoint.reset();
             if (Channel) {  // async actor doesn't hold channels, so channel is resumed in task runner actor
-                Channel->Resume();
+                Channel->ResumeByCheckpoint();
             }
         }
 
@@ -1609,37 +1607,42 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
     void InitializeTask() {
         for (ui32 i = 0; i < Task.InputsSize(); ++i) {
             const auto& inputDesc = Task.GetInputs(i);
+            auto watermarksMode = NDqProto::WATERMARKS_MODE_DISABLED;
             Y_ABORT_UNLESS(!inputDesc.HasSource() || inputDesc.ChannelsSize() == 0); // HasSource => no channels
 
-            if (inputDesc.HasTransform()) {
-                auto result = InputTransformsMap.emplace(
-                    i,
-                    TAsyncInputTransformHelper(LogPrefix, i, NDqProto::WATERMARKS_MODE_DISABLED)
-                );
-                YQL_ENSURE(result.second);
-            }
-
             if (inputDesc.HasSource()) {
-                const auto watermarksMode = inputDesc.GetSource().GetWatermarksMode();
+                watermarksMode = inputDesc.GetSource().GetWatermarksMode();
                 auto result = SourcesMap.emplace(
                     i,
                     static_cast<TDerived*>(this)->CreateInputHelper(LogPrefix, i, watermarksMode)
                 );
                 YQL_ENSURE(result.second);
             } else {
                 for (auto& channel : inputDesc.GetChannels()) {
+                    auto channelWatermarksMode = channel.GetWatermarksMode();
+                    if (channelWatermarksMode != NDqProto::EWatermarksMode::WATERMARKS_MODE_DISABLED) {
+                        watermarksMode = channelWatermarksMode;
+                    }
                     auto result = InputChannelsMap.emplace(
                         channel.GetId(),
                         TInputChannelInfo(
                             LogPrefix,
                             channel.GetId(),
                             channel.GetSrcStageId(),
-                            channel.GetWatermarksMode(),
+                            channelWatermarksMode,
                             channel.GetCheckpointingMode())
                     );
                     YQL_ENSURE(result.second);
                 }
             }
+
+            if (inputDesc.HasTransform()) {
+                auto result = InputTransformsMap.emplace(
+                    i,
+                    TAsyncInputTransformHelper(LogPrefix, i, watermarksMode)
+                );
+                YQL_ENSURE(result.second);
+            }
         }
 
         for (ui32 i = 0; i < Task.OutputsSize(); ++i) {
@@ -1694,12 +1697,6 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
                 WatermarksTracker.RegisterInputChannel(id);
             }
         }
-
-        for (const auto& [id, channel] : OutputChannelsMap) {
-            if (channel.WatermarksMode == NDqProto::EWatermarksMode::WATERMARKS_MODE_DEFAULT) {
-                WatermarksTracker.RegisterOutputChannel(id);
-            }
-        }
     }
 
     virtual const NYql::NDq::TDqTaskRunnerStats* GetTaskRunnerStats() = 0;