Refactor sink: shards state updating (#19071)

nikvas0 · nikvas0 · commit 3c4fc2c96a43 · 2025-06-02T15:11:36.000+03:00
diff --git a/ydb/core/kqp/common/kqp_tx_manager.cpp b/ydb/core/kqp/common/kqp_tx_manager.cpp
@@ -312,7 +312,7 @@ class TKqpTransactionManager : public IKqpTransactionManager {
 
     bool NeedCommit() const override {
         AFL_ENSURE(ActionsCount != 1 || IsSingleShard()); // ActionsCount == 1 then IsSingleShard()
-        const bool dontNeedCommit = IsEmpty() || IsReadOnly() && ((ActionsCount == 1) || HasSnapshot());
+        const bool dontNeedCommit = IsEmpty() || (IsReadOnly() && ((ActionsCount == 1) || HasSnapshot()));
         return !dontNeedCommit;
     }
 
diff --git a/ydb/core/kqp/runtime/kqp_write_actor.cpp b/ydb/core/kqp/runtime/kqp_write_actor.cpp
@@ -354,20 +354,13 @@ class TKqpTableWriteActor : public TActorBootstrapped<TKqpTableWriteActor> {
         ShardedWriteController->Close();
     }
 
-    void SetParentTraceId(NWilson::TTraceId traceId) {
-        ParentTraceId = std::move(traceId);
+    void CleanupClosedTokens() {
+        YQL_ENSURE(ShardedWriteController);
+        ShardedWriteController->CleanupClosedTokens();
     }
 
-    void UpdateShards() {
-        // TODO: Maybe there are better ways to initialize new shards...
-        for (const auto& shardInfo : ShardedWriteController->GetPendingShards()) {
-            TxManager->AddShard(shardInfo.ShardId, IsOlap, TablePath);
-            IKqpTransactionManager::TActionFlags flags = IKqpTransactionManager::EAction::WRITE;
-            if (shardInfo.HasRead) {
-                flags |= IKqpTransactionManager::EAction::READ;
-            }
-            TxManager->AddAction(shardInfo.ShardId, flags);
-        }
+    void SetParentTraceId(NWilson::TTraceId traceId) {
+        ParentTraceId = std::move(traceId);
     }
 
     bool IsClosed() const {
@@ -878,18 +871,30 @@ class TKqpTableWriteActor : public TActorBootstrapped<TKqpTableWriteActor> {
         }
     }
 
+    void UpdateShards() {
+        for (const auto& shardInfo : ShardedWriteController->ExtractShardUpdates()) {
+            TxManager->AddShard(shardInfo.ShardId, IsOlap, TablePath);
+            IKqpTransactionManager::TActionFlags flags = IKqpTransactionManager::EAction::WRITE;
+            if (shardInfo.HasRead) {
+                flags |= IKqpTransactionManager::EAction::READ;
+            }
+            TxManager->AddAction(shardInfo.ShardId, flags);
+        }
+    }
+
     void FlushBuffers() {
         ShardedWriteController->FlushBuffers();
         UpdateShards();
     }
 
-    bool Flush() {
-        for (const auto& shardInfo : ShardedWriteController->GetPendingShards()) {
-            if (!SendDataToShard(shardInfo.ShardId)) {
-                return false;
+    bool FlushToShards() {
+        bool ok = true;
+        ShardedWriteController->ForEachPendingShard([&](const auto& shardInfo) {
+            if (ok && !SendDataToShard(shardInfo.ShardId)) {
+                ok = false;
             }
-        }
-        return true;
+        });
+        return ok;
     }
 
     bool SendDataToShard(const ui64 shardId) {
@@ -1516,7 +1521,7 @@ class TKqpDirectWriteActor : public TActorBootstrapped<TKqpDirectWriteActor>, pu
             }
 
             if (Closed || outOfMemory) {
-                if (!WriteTableActor->Flush()) {
+                if (!WriteTableActor->FlushToShards()) {
                     return;
                 }
             }
@@ -1933,7 +1938,7 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
             CA_LOG_D("Flush data");
             for (auto& [_, info] : WriteInfos) {
                 if (info.WriteTableActor->IsReady()) {
-                    if (!info.WriteTableActor->Flush()) {
+                    if (!info.WriteTableActor->FlushToShards()) {
                         return false;
                     }
                 }
@@ -2807,6 +2812,7 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
         Y_ABORT_UNLESS(GetTotalMemory() == 0);
 
         for (auto& [_, info] : WriteInfos) {
+            info.WriteTableActor->CleanupClosedTokens();
             info.WriteTableActor->Unlink();
         }
     }
diff --git a/ydb/core/kqp/runtime/kqp_write_table.cpp b/ydb/core/kqp/runtime/kqp_write_table.cpp
@@ -1187,17 +1187,15 @@ class TShardsInfo {
         return insertIt->second;
     }
 
-    TVector<IShardedWriteController::TPendingShardInfo> GetPendingShards() const {
-        TVector<IShardedWriteController::TPendingShardInfo> result;
+    void ForEachPendingShard(std::function<void(const IShardedWriteController::TPendingShardInfo&)>&& callback) const {
         for (const auto& [id, shard] : ShardsInfo) {
             if (!shard.IsEmpty() && shard.GetSendAttempts() == 0) {
-                result.push_back(IShardedWriteController::TPendingShardInfo{
+                callback(IShardedWriteController::TPendingShardInfo{
                     .ShardId = id,
                     .HasRead = shard.HasRead(),
                 });
             }
         }
-        return result;
     }
 
     bool Has(ui64 shardId) const {
@@ -1382,6 +1380,18 @@ class TShardedWriteController : public IShardedWriteController {
         }
     }
 
+    void CleanupClosedTokens() override {
+        AFL_ENSURE(IsEmpty());
+        for (auto it = WriteInfos.begin(); it != WriteInfos.end();) {
+            if (it->second.Closed) {
+                AFL_ENSURE(it->second.Serializer->IsFinished());
+                it = WriteInfos.erase(it);
+            } else {
+                ++it;
+            }
+        }
+    }
+
     void FlushBuffers() override {
         TVector<TWriteToken> writeTokensFoFlush;
         for (const auto& [token, writeInfo] : WriteInfos) {
@@ -1420,8 +1430,14 @@ class TShardedWriteController : public IShardedWriteController {
         }
     }
 
-    TVector<TPendingShardInfo> GetPendingShards() const override {
-        return ShardsInfo.GetPendingShards();
+    void ForEachPendingShard(std::function<void(const TPendingShardInfo&)>&& callback) const override {
+        ShardsInfo.ForEachPendingShard(std::move(callback));
+    }
+
+    std::vector<TPendingShardInfo> ExtractShardUpdates() override {
+        std::vector<TPendingShardInfo> shardUpdates;
+        std::swap(shardUpdates, ShardUpdates);
+        return shardUpdates;
     }
 
     TVector<ui64> GetShardsIds() const override {
@@ -1580,11 +1596,16 @@ class TShardedWriteController : public IShardedWriteController {
             for (auto& [shardId, batches] : writeInfo.Serializer->FlushBatchesForce()) {
                 for (auto& batch : batches) {
                     if (batch && !batch->IsEmpty()) {
+                        const bool hasRead = (writeInfo.Metadata.OperationType == NKikimrDataEvents::TEvWrite::TOperation::OPERATION_INSERT
+                                || writeInfo.Metadata.OperationType == NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPDATE);
                         ShardsInfo.GetShard(shardId).PushBatch(TBatchWithMetadata{
                             .Token = token,
                             .Data = std::move(batch),
-                            .HasRead = (writeInfo.Metadata.OperationType == NKikimrDataEvents::TEvWrite::TOperation::OPERATION_INSERT
-                                || writeInfo.Metadata.OperationType == NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPDATE),
+                            .HasRead = hasRead,
+                        });
+                        ShardUpdates.push_back(IShardedWriteController::TPendingShardInfo{
+                            .ShardId = shardId,
+                            .HasRead = hasRead,
                         });
                     }
                 }
@@ -1597,11 +1618,16 @@ class TShardedWriteController : public IShardedWriteController {
                     if (!batch || batch->IsEmpty()) {
                         break;
                     }
+                    const bool hasRead = (writeInfo.Metadata.OperationType == NKikimrDataEvents::TEvWrite::TOperation::OPERATION_INSERT
+                        || writeInfo.Metadata.OperationType == NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPDATE);
                     shard.PushBatch(TBatchWithMetadata{
                         .Token = token,
                         .Data = std::move(batch),
-                        .HasRead = (writeInfo.Metadata.OperationType == NKikimrDataEvents::TEvWrite::TOperation::OPERATION_INSERT
-                                || writeInfo.Metadata.OperationType == NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPDATE),
+                        .HasRead = hasRead,
+                    });
+                    ShardUpdates.push_back(IShardedWriteController::TPendingShardInfo{
+                        .ShardId = shardId,
+                        .HasRead = hasRead,
                     });
                 }
             }
@@ -1615,6 +1641,7 @@ class TShardedWriteController : public IShardedWriteController {
                 shard.MakeNextBatches(1);
             } else {
                 shard.MakeNextBatches(std::nullopt);
+                AFL_ENSURE(shard.GetBatchesInFlight() == shard.Size());
             }
         }
     }
@@ -1646,6 +1673,7 @@ class TShardedWriteController : public IShardedWriteController {
     TWriteToken CurrentWriteToken = 0;
 
     TShardsInfo ShardsInfo;
+    std::vector<IShardedWriteController::TPendingShardInfo> ShardUpdates;
 
     std::optional<NSchemeCache::TSchemeCacheNavigate::TEntry> SchemeEntry;
     std::shared_ptr<const TVector<TKeyDesc::TPartitionInfo>> Partitioning;
diff --git a/ydb/core/kqp/runtime/kqp_write_table.h b/ydb/core/kqp/runtime/kqp_write_table.h
@@ -69,6 +69,8 @@ class IShardedWriteController : public TThrRefBase {
     virtual void Write(TWriteToken token, IDataBatchPtr&& data) = 0;
     virtual void Close(TWriteToken token) = 0;
 
+    virtual void CleanupClosedTokens() = 0;
+
     virtual void FlushBuffers() = 0;
 
     virtual void Close() = 0;
@@ -79,7 +81,9 @@ class IShardedWriteController : public TThrRefBase {
         ui64 ShardId;
         bool HasRead;
     };
-    virtual TVector<TPendingShardInfo> GetPendingShards() const = 0;
+    virtual void ForEachPendingShard(std::function<void(const TPendingShardInfo&)>&& callback) const = 0;
+    virtual std::vector<TPendingShardInfo> ExtractShardUpdates() = 0;
+
     virtual ui64 GetShardsCount() const = 0;
     virtual TVector<ui64> GetShardsIds() const = 0;
 

Original file line number	Diff line number	Diff line change
`@@ -312,7 +312,7 @@ class TKqpTransactionManager : public IKqpTransactionManager {`
`312`	`312`
`313`	`313`	`bool NeedCommit() const override {`
`314`	`314`	`AFL_ENSURE(ActionsCount != 1 \|\| IsSingleShard()); // ActionsCount == 1 then IsSingleShard()`
`315`		`- const bool dontNeedCommit = IsEmpty() \|\| IsReadOnly() && ((ActionsCount == 1) \|\| HasSnapshot());`
	`315`	`+ const bool dontNeedCommit = IsEmpty() \|\| (IsReadOnly() && ((ActionsCount == 1) \|\| HasSnapshot()));`
`316`	`316`	`return !dontNeedCommit;`
`317`	`317`	`}`
`318`	`318`
Original file line number	Diff line number	Diff line change
`@@ -354,20 +354,13 @@ class TKqpTableWriteActor : public TActorBootstrapped<TKqpTableWriteActor> {`
`354`	`354`	`ShardedWriteController->Close();`
`355`	`355`	`}`
`356`	`356`
`357`		`- void SetParentTraceId(NWilson::TTraceId traceId) {`
`358`		`- ParentTraceId = std::move(traceId);`
	`357`	`+ void CleanupClosedTokens() {`
	`358`	`+ YQL_ENSURE(ShardedWriteController);`
	`359`	`+ ShardedWriteController->CleanupClosedTokens();`
`359`	`360`	`}`
`360`	`361`
`361`		`- void UpdateShards() {`
`362`		`- // TODO: Maybe there are better ways to initialize new shards...`
`363`		`- for (const auto& shardInfo : ShardedWriteController->GetPendingShards()) {`
`364`		`- TxManager->AddShard(shardInfo.ShardId, IsOlap, TablePath);`
`365`		`- IKqpTransactionManager::TActionFlags flags = IKqpTransactionManager::EAction::WRITE;`
`366`		`- if (shardInfo.HasRead) {`
`367`		`- flags \|= IKqpTransactionManager::EAction::READ;`
`368`		`- }`
`369`		`- TxManager->AddAction(shardInfo.ShardId, flags);`
`370`		`- }`
	`362`	`+ void SetParentTraceId(NWilson::TTraceId traceId) {`
	`363`	`+ ParentTraceId = std::move(traceId);`
`371`	`364`	`}`
`372`	`365`
`373`	`366`	`bool IsClosed() const {`
`@@ -878,18 +871,30 @@ class TKqpTableWriteActor : public TActorBootstrapped<TKqpTableWriteActor> {`
`878`	`871`	`}`
`879`	`872`	`}`
`880`	`873`
	`874`	`+ void UpdateShards() {`
	`875`	`+ for (const auto& shardInfo : ShardedWriteController->ExtractShardUpdates()) {`
	`876`	`+ TxManager->AddShard(shardInfo.ShardId, IsOlap, TablePath);`
	`877`	`+ IKqpTransactionManager::TActionFlags flags = IKqpTransactionManager::EAction::WRITE;`
	`878`	`+ if (shardInfo.HasRead) {`
	`879`	`+ flags \|= IKqpTransactionManager::EAction::READ;`
	`880`	`+ }`
	`881`	`+ TxManager->AddAction(shardInfo.ShardId, flags);`
	`882`	`+ }`
	`883`	`+ }`
	`884`	`+`
`881`	`885`	`void FlushBuffers() {`
`882`	`886`	`ShardedWriteController->FlushBuffers();`
`883`	`887`	`UpdateShards();`
`884`	`888`	`}`
`885`	`889`
`886`		`- bool Flush() {`
`887`		`- for (const auto& shardInfo : ShardedWriteController->GetPendingShards()) {`
`888`		`- if (!SendDataToShard(shardInfo.ShardId)) {`
`889`		`- return false;`
	`890`	`+ bool FlushToShards() {`
	`891`	`+ bool ok = true;`
	`892`	`+ ShardedWriteController->ForEachPendingShard([&](const auto& shardInfo) {`
	`893`	`+ if (ok && !SendDataToShard(shardInfo.ShardId)) {`
	`894`	`+ ok = false;`
`890`	`895`	`}`
`891`		`- }`
`892`		`- return true;`
	`896`	`+ });`
	`897`	`+ return ok;`
`893`	`898`	`}`
`894`	`899`
`895`	`900`	`bool SendDataToShard(const ui64 shardId) {`
`@@ -1516,7 +1521,7 @@ class TKqpDirectWriteActor : public TActorBootstrapped<TKqpDirectWriteActor>, pu`
`1516`	`1521`	`}`
`1517`	`1522`
`1518`	`1523`	`if (Closed \|\| outOfMemory) {`
`1519`		`- if (!WriteTableActor->Flush()) {`
	`1524`	`+ if (!WriteTableActor->FlushToShards()) {`
`1520`	`1525`	`return;`
`1521`	`1526`	`}`
`1522`	`1527`	`}`
`@@ -1933,7 +1938,7 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub`
`1933`	`1938`	`CA_LOG_D("Flush data");`
`1934`	`1939`	`for (auto& [_, info] : WriteInfos) {`
`1935`	`1940`	`if (info.WriteTableActor->IsReady()) {`
`1936`		`- if (!info.WriteTableActor->Flush()) {`
	`1941`	`+ if (!info.WriteTableActor->FlushToShards()) {`
`1937`	`1942`	`return false;`
`1938`	`1943`	`}`
`1939`	`1944`	`}`
`@@ -2807,6 +2812,7 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub`
`2807`	`2812`	`Y_ABORT_UNLESS(GetTotalMemory() == 0);`
`2808`	`2813`
`2809`	`2814`	`for (auto& [_, info] : WriteInfos) {`
	`2815`	`+ info.WriteTableActor->CleanupClosedTokens();`
`2810`	`2816`	`info.WriteTableActor->Unlink();`
`2811`	`2817`	`}`
`2812`	`2818`	`}`