Merge c33c319 into fcc71af

Author: zverevgeny

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

@@ -313,17 +313,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
     virtual void DoExecuteImpl() = 0;
     virtual void DoTerminateImpl() {}
 
-    virtual bool DoHandleChannelsAfterFinishImpl() {
-        Y_ABORT_UNLESS(Checkpoints);
-
-        if (Checkpoints->HasPendingCheckpoint() && !Checkpoints->ComputeActorStateSaved() && ReadyToCheckpoint()) {
-            Checkpoints->DoCheckpoint();
-        }
-
-        // Send checkpoints to output channels.
-        ProcessOutputsImpl(ERunStatus::Finished);
-        return true;  // returns true, when channels were handled synchronously
-    }
+    virtual bool DoHandleChannelsAfterFinishImpl() = 0;
 
     void ProcessOutputsImpl(ERunStatus status) {
         ProcessOutputsState.LastRunStatus = status;
@@ -621,131 +611,50 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
         }
     }
 
-public:
-    i64 GetInputChannelFreeSpace(ui64 channelId) const override {
-        const TInputChannelInfo* inputChannel = InputChannelsMap.FindPtr(channelId);
-        YQL_ENSURE(inputChannel, "task: " << Task.GetId() << ", unknown input channelId: " << channelId);
-
-        return inputChannel->Channel->GetFreeSpace();
-    }
-
-    void TakeInputChannelData(TChannelDataOOB&& channelData, bool ack) override {
-        TInputChannelInfo* inputChannel = InputChannelsMap.FindPtr(channelData.Proto.GetChannelId());
-        YQL_ENSURE(inputChannel, "task: " << Task.GetId() << ", unknown input channelId: " << channelData.Proto.GetChannelId());
-
-        auto channel = inputChannel->Channel;
-
-        if (channelData.RowCount()) {
-            TDqSerializedBatch batch;
-            batch.Proto = std::move(*channelData.Proto.MutableData());
-            batch.Payload = std::move(channelData.Payload);
-            auto guard = BindAllocator();
-            channel->Push(std::move(batch));
-        }
-
-        if (channelData.Proto.HasCheckpoint()) {
-            Y_ABORT_UNLESS(inputChannel->CheckpointingMode != NDqProto::CHECKPOINTING_MODE_DISABLED);
-            Y_ABORT_UNLESS(Checkpoints);
-            const auto& checkpoint = channelData.Proto.GetCheckpoint();
-            inputChannel->Pause(checkpoint);
-            Checkpoints->RegisterCheckpoint(checkpoint, channelData.Proto.GetChannelId());
-        }
-
-        if (channelData.Proto.GetFinished()) {
-            channel->Finish();
-        }
-
-        if (ack) {
-            Channels->SendChannelDataAck(channel->GetChannelId(), channel->GetFreeSpace());
-        }
-
-        ContinueExecute(EResumeSource::CATakeInput);
-    }
-
-    void PeerFinished(ui64 channelId) override {
-        TOutputChannelInfo* outputChannel = OutputChannelsMap.FindPtr(channelId);
-        YQL_ENSURE(outputChannel, "task: " << Task.GetId() << ", output channelId: " << channelId);
-
-        outputChannel->Finished = true;
-        outputChannel->Channel->Finish();
-
-        CA_LOG_D("task: " << Task.GetId() << ", output channelId: " << channelId << " finished prematurely, "
-            << " about to clear buffer");
-
-        {
-            auto guard = BindAllocator();
-            ui32 dropRows = outputChannel->Channel->Drop();
+protected: //TDqComputeActorChannels::ICallbacks
+    //i64 GetInputChannelFreeSpace(ui64 channelId) is pure and must be overridded in derived class
 
-            CA_LOG_I("task: " << Task.GetId() << ", output channelId: " << channelId << " finished prematurely, "
-                << "drop " << dropRows << " rows");
-        }
+    //void TakeInputChannelData(TChannelDataOOB&& channelData, bool ack) is pure and must be overridded in derived class
 
-        DoExecute();
-    }
+    // void PeerFinished(ui64 channelId) is pure and must be overridded in derived class
 
-    void ResumeExecution(EResumeSource source) override {
+    void ResumeExecution(EResumeSource source) override final{
         ContinueExecute(source);
     }
 
-    void OnSinkStateSaved(NDqProto::TSinkState&& state, ui64 outputIndex, const NDqProto::TCheckpoint& checkpoint) override {
+protected:
+    void OnSinkStateSaved(NDqProto::TSinkState&& state, ui64 outputIndex, const NDqProto::TCheckpoint& checkpoint) override final {
         Y_ABORT_UNLESS(Checkpoints); // If we are checkpointing, we must have already constructed "checkpoints" object.
         Checkpoints->OnSinkStateSaved(std::move(state), outputIndex, checkpoint);
     }
 
-    void OnTransformStateSaved(NDqProto::TSinkState&& state, ui64 outputIndex, const NDqProto::TCheckpoint& checkpoint) override {
+    void OnTransformStateSaved(NDqProto::TSinkState&& state, ui64 outputIndex, const NDqProto::TCheckpoint& checkpoint) override final {
         Y_ABORT_UNLESS(Checkpoints); // If we are checkpointing, we must have already constructed "checkpoints" object.
         Checkpoints->OnTransformStateSaved(std::move(state), outputIndex, checkpoint);
     }
 
-    void OnSinkFinished(ui64 outputIndex) override {
+    void OnSinkFinished(ui64 outputIndex) override final {
         SinksMap.at(outputIndex).FinishIsAcknowledged = true;
         ContinueExecute(EResumeSource::CASinkFinished);
     }
 
-    void OnTransformFinished(ui64 outputIndex) override {
+    void OnTransformFinished(ui64 outputIndex) override final {
         OutputTransformsMap.at(outputIndex).FinishIsAcknowledged = true;
         ContinueExecute(EResumeSource::CATransformFinished);
     }
+    
+protected: //TDqComputeActorCheckpoints::ICallbacks
+    //bool ReadyToCheckpoint() is pure and must be overriden in a derived class
 
-protected:
-    bool ReadyToCheckpoint() const override {
-        for (auto& [id, channelInfo] : InputChannelsMap) {
-            if (channelInfo.CheckpointingMode == NDqProto::CHECKPOINTING_MODE_DISABLED) {
-                continue;
-            }
-
-            if (!channelInfo.IsPaused()) {
-                return false;
-            }
-            if (!channelInfo.Channel->Empty()) {
-                return false;
-            }
-        }
-        return true;
-    }
-
-    void CommitState(const NDqProto::TCheckpoint& checkpoint) override {
+    void CommitState(const NDqProto::TCheckpoint& checkpoint) override final{
         CA_LOG_D("Commit state");
         for (auto& [inputIndex, source] : SourcesMap) {
             Y_ABORT_UNLESS(source.AsyncInput);
             source.AsyncInput->CommitState(checkpoint);
         }
     }
 
-    void InjectBarrierToOutputs(const NDqProto::TCheckpoint& checkpoint) override {
-        Y_ABORT_UNLESS(CheckpointingMode != NDqProto::CHECKPOINTING_MODE_DISABLED);
-        for (const auto& [id, channelInfo] : OutputChannelsMap) {
-            if (!channelInfo.IsTransformOutput) {
-                channelInfo.Channel->Push(NDqProto::TCheckpoint(checkpoint));
-            }
-        }
-        for (const auto& [outputIndex, sink] : SinksMap) {
-            sink.Buffer->Push(NDqProto::TCheckpoint(checkpoint));
-        }
-        for (const auto& [outputIndex, transform] : OutputTransformsMap) {
-            transform.Buffer->Push(NDqProto::TCheckpoint(checkpoint));
-        }
-    }
+    // void InjectBarrierToOutputs(const NDqProto::TCheckpoint& checkpoint) is pure and must be overriden in a derived class
 
     void ResumeInputsByWatermark(TInstant watermark) {
         for (auto& [id, sourceInfo] : SourcesMap) {
@@ -771,15 +680,16 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
         }
     }
 
-    void ResumeInputsByCheckpoint() override {
+    void ResumeInputsByCheckpoint() override final {
         for (auto& [id, channelInfo] : InputChannelsMap) {
             channelInfo.ResumeByCheckpoint();
         }
     }
 
+protected:
     virtual void DoLoadRunnerState(TString&& blob) = 0;
 
-    void LoadState(NDqProto::TComputeActorState&& state) override {
+    void LoadState(NDqProto::TComputeActorState&& state) override final {
         CA_LOG_D("Load state");
         TMaybe<TString> error = Nothing();
         const NDqProto::TMiniKqlProgramState& mkqlProgramState = state.GetMiniKqlProgram();
@@ -814,13 +724,13 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
         }
     }
 
-    void Start() override {
+    void Start() override final {
         Running = true;
         State = NDqProto::COMPUTE_STATE_EXECUTING;
         ContinueExecute(EResumeSource::CAStart);
     }
 
-    void Stop() override {
+    void Stop() override final {
         Running = false;
         State = NDqProto::COMPUTE_STATE_UNKNOWN;
     }
@@ -1255,99 +1165,12 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
     NKikimr::NMiniKQL::TScopedAlloc& GetAllocator() {
         return *Alloc.get();
     }
-private:
-    virtual const TDqMemoryQuota::TProfileStats* GetMemoryProfileStats() const {
-        Y_ABORT_UNLESS(MemoryQuota);
-        return MemoryQuota->GetProfileStats();
-    }
-
-    virtual void DrainOutputChannel(TOutputChannelInfo& outputChannel) {
-        YQL_ENSURE(!outputChannel.Finished || Checkpoints);
-
-        const bool wasFinished = outputChannel.Finished;
-        auto channelId = outputChannel.Channel->GetChannelId();
 
-        CA_LOG_T("About to drain channelId: " << channelId
-            << ", hasPeer: " << outputChannel.HasPeer
-            << ", finished: " << outputChannel.Channel->IsFinished());
+    virtual const TDqMemoryQuota::TProfileStats* GetMemoryProfileStats() const = 0;
 
-        ProcessOutputsState.HasDataToSend |= !outputChannel.Finished;
-        ProcessOutputsState.AllOutputsFinished &= outputChannel.Finished;
+    virtual void DrainOutputChannel(TOutputChannelInfo& outputChannel) = 0;
 
-        UpdateBlocked(outputChannel, !Channels->HasFreeMemoryInChannel(channelId));
-
-        ui32 sentChunks = 0;
-        while ((!outputChannel.Finished || Checkpoints) &&
-            Channels->HasFreeMemoryInChannel(outputChannel.ChannelId))
-        {
-            const static ui32 drainPackSize = 16;
-            std::vector<typename TOutputChannelInfo::TDrainedChannelMessage> channelData = outputChannel.DrainChannel(drainPackSize);
-            ui32 idx = 0;
-            for (auto&& i : channelData) {
-                if (auto* w = i.GetWatermarkOptional()) {
-                    CA_LOG_I("Resume inputs by watermark");
-                    // This is excessive, inputs should be resumed after async CA received response with watermark from task runner.
-                    // But, let it be here, it's better to have the same code as in checkpoints
-                    ResumeInputsByWatermark(TInstant::MicroSeconds(w->GetTimestampUs()));
-                }
-                if (i.GetCheckpointOptional()) {
-                    CA_LOG_I("Resume inputs by checkpoint");
-                    ResumeInputsByCheckpoint();
-                }
-
-                Channels->SendChannelData(i.BuildChannelData(outputChannel.ChannelId), ++idx == channelData.size());
-                ++sentChunks;
-            }
-            if (drainPackSize != channelData.size()) {
-                if (!outputChannel.Finished) {
-                    CA_LOG_T("output channelId: " << outputChannel.ChannelId << ", nothing to send and is not finished");
-                }
-                break;
-            }
-        }
-
-        ProcessOutputsState.HasDataToSend |= !outputChannel.Finished;
-        ProcessOutputsState.AllOutputsFinished &= outputChannel.Finished;
-        ProcessOutputsState.DataWasSent |= (!wasFinished && outputChannel.Finished) || sentChunks;
-    }
-
-    virtual void DrainAsyncOutput(ui64 outputIndex, TAsyncOutputInfoBase& outputInfo) {
-        ProcessOutputsState.AllOutputsFinished &= outputInfo.Finished;
-        if (outputInfo.Finished && !Checkpoints) {
-            return;
-        }
-
-        Y_ABORT_UNLESS(outputInfo.Buffer);
-        Y_ABORT_UNLESS(outputInfo.AsyncOutput);
-        Y_ABORT_UNLESS(outputInfo.Actor);
-
-        const ui32 allowedOvercommit = AllowedChannelsOvercommit();
-        const i64 sinkFreeSpaceBeforeSend = outputInfo.AsyncOutput->GetFreeSpace();
-
-        i64 toSend = sinkFreeSpaceBeforeSend + allowedOvercommit;
-        CA_LOG_D("About to drain async output " << outputIndex
-            << ". FreeSpace: " << sinkFreeSpaceBeforeSend
-            << ", allowedOvercommit: " << allowedOvercommit
-            << ", toSend: " << toSend
-            << ", finished: " << outputInfo.Buffer->IsFinished());
-
-        i64 sent = 0;
-        while (toSend > 0 && (!outputInfo.Finished || Checkpoints)) {
-            const ui32 sentChunk = SendDataChunkToAsyncOutput(outputIndex, outputInfo, toSend);
-            if (sentChunk == 0) {
-                break;
-            }
-            sent += sentChunk;
-            toSend = outputInfo.AsyncOutput->GetFreeSpace() + allowedOvercommit;
-        }
-
-        CA_LOG_D("Drain async output " << outputIndex
-            << ". Free space decreased: " << (sinkFreeSpaceBeforeSend - outputInfo.AsyncOutput->GetFreeSpace())
-            << ", sent data from buffer: " << sent);
-
-        ProcessOutputsState.HasDataToSend |= !outputInfo.Finished;
-        ProcessOutputsState.DataWasSent |= outputInfo.Finished || sent;
-    }
+    virtual void DrainAsyncOutput(ui64 outputIndex, TAsyncOutputInfoBase& outputInfo) = 0;
 
     ui32 SendDataChunkToAsyncOutput(ui64 outputIndex, TAsyncOutputInfoBase& outputInfo, ui64 bytes) {
         auto sink = outputInfo.Buffer;
@@ -1599,7 +1422,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
         InternalError(fatalCode, issues);
     }
 
-    void OnSinkError(ui64 outputIndex, const TIssues& issues, NYql::NDqProto::StatusIds::StatusCode fatalCode) override {
+    void OnSinkError(ui64 outputIndex, const TIssues& issues, NYql::NDqProto::StatusIds::StatusCode fatalCode) override final {
         if (fatalCode == NYql::NDqProto::StatusIds::UNSPECIFIED) {
             SinksMap.at(outputIndex).IssuesBuffer.Push(issues);
             return;
@@ -1609,7 +1432,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
         InternalError(fatalCode, issues);
     }
 
-    void OnOutputTransformError(ui64 outputIndex, const TIssues& issues, NYql::NDqProto::StatusIds::StatusCode fatalCode) override {
+    void OnOutputTransformError(ui64 outputIndex, const TIssues& issues, NYql::NDqProto::StatusIds::StatusCode fatalCode) override final {
         if (fatalCode == NYql::NDqProto::StatusIds::UNSPECIFIED) {
             OutputTransformsMap.at(outputIndex).IssuesBuffer.Push(issues);
             return;
@@ -1635,7 +1458,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
         return true;
     }
 
-    virtual ui64 CalcMkqlMemoryLimit() {
+    ui64 CalcMkqlMemoryLimit() {
         auto& opts = Task.GetProgram().GetSettings();
         return opts.GetHasMapJoin()/* || opts.GetHasSort()*/
             ? MemoryLimits.MkqlHeavyProgramMemoryLimit
@@ -1738,9 +1561,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
     virtual const NYql::NDq::TTaskRunnerStatsBase* GetTaskRunnerStats() = 0;
     virtual const NYql::NDq::TDqMeteringStats* GetMeteringStats() = 0;
 
-    virtual const IDqAsyncOutputBuffer* GetSink(ui64, const TAsyncOutputInfoBase& sinkInfo) const {
-        return sinkInfo.Buffer.Get();
-    }
+    virtual const IDqAsyncOutputBuffer* GetSink(ui64 outputIdx, const TAsyncOutputInfoBase& sinkInfo) const = 0;
 
 public: