Add read iterator cancellation to ReadRows RPC (#19376)

SammyVimes · SammyVimes · commit 7ddf958a8d78 · 2025-06-17T17:54:17.000+03:00
(cherry picked from commit 4ee8a57)
diff --git a/ydb/core/grpc_services/rpc_read_rows.cpp b/ydb/core/grpc_services/rpc_read_rows.cpp
@@ -40,6 +40,7 @@ struct RequestedKeyColumn {
 struct TShardReadState {
     std::vector<TOwnedCellVec> Keys;
     ui32 FirstUnprocessedQuery = 0;
+    Ydb::StatusIds::StatusCode Status = Ydb::StatusIds::STATUS_CODE_UNSPECIFIED;
 };
 
 }
@@ -338,6 +339,14 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
         Become(&TThis::MainState);
 
         LOG_DEBUG_S(TlsActivationContext->AsActorContext(), NKikimrServices::RPC_REQUEST, "TReadRowsRPC bootstraped ");
+
+        auto selfId = ctx.SelfID;
+        auto* actorSystem = ctx.ActorSystem();
+        auto clientLostCb = [selfId, actorSystem]() {
+            actorSystem->Send(selfId, new TRpcServices::TEvForgetOperation());
+        };
+
+        Request->SetFinishAction(std::move(clientLostCb));
     }
 
     bool ResolveTable() {
@@ -454,7 +463,7 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
         }
     }
 
-    void Handle(TEvTxProxySchemeCache::TEvResolveKeySetResult::TPtr &ev) {
+    void Handle(TEvTxProxySchemeCache::TEvResolveKeySetResult::TPtr& ev) {
         TEvTxProxySchemeCache::TEvResolveKeySetResult *msg = ev->Get();
         auto& resolvePartitionsResult = msg->Request;
 
@@ -519,8 +528,7 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
             if (it == ShardIdToReadState.end()) {
                 TStringStream ss;
                 ss << "Got unknown shardId from TEvReadResult# " << shardId << ", status# " << statusCode;
-                ReplyWithError(statusCode, ss.Str(), &issues);
-                return;
+                return ReplyWithError(statusCode, ss.Str(), &issues);
             }
 
             switch (statusCode) {
@@ -531,6 +539,7 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
                 if (Retries < MaxTotalRetries) {
                     TStringStream ss;
                     ss << "Reached MaxRetries count for DataShard# " << shardId << ", status# " << statusCode;
+                    it->second.Status = statusCode;
                     ReplyWithError(statusCode, ss.Str(), &issues);
                 } else {
                     SendRead(shardId, it->second);
@@ -545,8 +554,8 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
                 if (statusCode != Ydb::StatusIds::OVERLOADED) {
                     statusCode = Ydb::StatusIds::ABORTED;
                 }
-                ReplyWithError(statusCode, ss.Str(), &issues);
-                return;
+                it->second.Status = statusCode;
+                return ReplyWithError(statusCode, ss.Str(), &issues);
             }
             }
             if (!msg->Record.HasFinished() || !msg->Record.GetFinished()) {
@@ -562,6 +571,9 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
                 // we just wait for the next batch of results.
                 it->second.FirstUnprocessedQuery = token.GetFirstUnprocessedQuery();
                 ReadsInFlight++;
+            } else {
+                // Read for this shard has finished
+                it->second.Status = statusCode;
             }
         }
         LOG_DEBUG_S(TlsActivationContext->AsActorContext(), NKikimrServices::RPC_REQUEST, "TReadRowsRPC TEvReadResult RowsCount: " << msg->GetRowsCount());
@@ -720,14 +732,48 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
         PassAway();
     }
 
+    void CancelReads() {
+        TStringStream ss;
+        ss << "TReadRowsRPC CancelReads, shardIds# [";
+
+        bool hasActiveReads = false;
+
+        for (const auto& [shardId, state] : ShardIdToReadState) {
+            if (state.Status != Ydb::StatusIds::STATUS_CODE_UNSPECIFIED) {
+                // Read has already finished for this shard
+                continue;
+            }
+            auto request = std::make_unique<TEvDataShard::TEvReadCancel>();
+            auto& record = request->Record;
+            record.SetReadId(shardId); // shardId is also a readId
+            Send(PipeCache, new TEvPipeCache::TEvForward(request.release(), shardId, true), IEventHandle::FlagTrackDelivery, 0, Span.GetTraceId());
+            ss << shardId << ", ";
+            hasActiveReads = true;
+        }
+
+        ss << "]";
+
+        if (hasActiveReads) {
+            LOG_WARN_S(TlsActivationContext->AsActorContext(), NKikimrServices::RPC_REQUEST, ss.Str());
+        }
+    }
+
     void HandleTimeout(TEvents::TEvWakeup::TPtr&) {
-        return ReplyWithError(Ydb::StatusIds::TIMEOUT, TStringBuilder() << "ReadRows from table " << GetTable()
+        ReplyWithError(Ydb::StatusIds::TIMEOUT, TStringBuilder() << "ReadRows from table " << GetTable()
             << " timed out, duration: " << (TAppData::TimeProvider->Now() - StartTime).Seconds() << " sec");
     }
 
+    void HandleForget(TRpcServices::TEvForgetOperation::TPtr& ev) {
+        Y_UNUSED(ev);
+
+        ReplyWithError(Ydb::StatusIds::CANCELLED, TStringBuilder() << "ReadRows from table " << GetTable()
+            << " cancelled, because client disconnected");
+    }
+
     void ReplyWithError(const Ydb::StatusIds::StatusCode& status, const TString& errorMsg,
         const ::google::protobuf::RepeatedPtrField<Ydb::Issue::IssueMessage>* issues = nullptr)
     {
+        CancelReads();
         LOG_ERROR_S(TlsActivationContext->AsActorContext(), NKikimrServices::RPC_REQUEST, "TReadRowsRPC ReplyWithError: " << errorMsg);
         SendResult(status, errorMsg, issues);
     }
@@ -736,7 +782,7 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
         return ReplyWithError(Ydb::StatusIds::INTERNAL_ERROR, "Internal error: pipe cache is not available, the cluster might not be configured properly");
     }
 
-    void Handle(TEvPipeCache::TEvDeliveryProblem::TPtr &ev) {
+    void Handle(TEvPipeCache::TEvDeliveryProblem::TPtr& ev) {
         return ReplyWithError(Ydb::StatusIds::UNAVAILABLE, TStringBuilder() << "Failed to connect to shard " << ev->Get()->TabletId);
     }
 
@@ -760,6 +806,7 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
             hFunc(TEvPipeCache::TEvDeliveryProblem, Handle);
 
             hFunc(TEvents::TEvWakeup, HandleTimeout);
+            hFunc(TRpcServices::TEvForgetOperation, HandleForget);
         }
     }
 
diff --git a/ydb/core/kqp/ut/common/kqp_ut_common.cpp b/ydb/core/kqp/ut/common/kqp_ut_common.cpp
@@ -155,7 +155,15 @@ TKikimrRunner::TKikimrRunner(const TKikimrSettings& settings) {
     }
 
     Server.Reset(MakeHolder<Tests::TServer>(*ServerSettings));
-    Server->EnableGRpc(grpcPort);
+
+    if (settings.GrpcServerOptions) {
+        auto options = settings.GrpcServerOptions;
+        options->SetPort(grpcPort);
+        options->SetHost("localhost");
+        Server->EnableGRpc(*options);
+    } else {
+        Server->EnableGRpc(grpcPort);
+    }
 
     RunCall([this, domain = settings.DomainRoot] {
         this->Server->SetupDefaultProfiles();
diff --git a/ydb/core/kqp/ut/common/kqp_ut_common.h b/ydb/core/kqp/ut/common/kqp_ut_common.h
@@ -90,6 +90,7 @@ struct TKikimrSettings: public TTestFeatureFlagsHolder<TKikimrSettings> {
     NMonitoring::TDynamicCounterPtr CountersRoot = MakeIntrusive<NMonitoring::TDynamicCounters>();
     std::shared_ptr<NYql::NDq::IS3ActorsFactory> S3ActorsFactory = NYql::NDq::CreateDefaultS3ActorsFactory();
     NKikimrConfig::TImmediateControlsConfig Controls;
+    TMaybe<NYdbGrpc::TServerOptions> GrpcServerOptions;
 
     TKikimrSettings()
     {
@@ -128,6 +129,7 @@ struct TKikimrSettings: public TTestFeatureFlagsHolder<TKikimrSettings> {
             return *this;
     }
     TKikimrSettings& SetControls(const NKikimrConfig::TImmediateControlsConfig& value) { Controls = value; return *this; }
+    TKikimrSettings& SetGrpcServerOptions(const NYdbGrpc::TServerOptions& grpcServerOptions) { GrpcServerOptions = grpcServerOptions; return *this; };
 };
 
 class TKikimrRunner {
diff --git a/ydb/core/kqp/ut/opt/kqp_kv_ut.cpp b/ydb/core/kqp/ut/opt/kqp_kv_ut.cpp
@@ -1,5 +1,6 @@
 #include <ydb/core/kqp/ut/common/kqp_ut_common.h>
 
+#include <ydb/core/base/blobstorage.h>
 #include <ydb/core/tx/datashard/datashard.h>
 #include <yql/essentials/parser/pg_catalog/catalog.h>
 #include <yql/essentials/parser/pg_wrapper/interface/codec.h>
@@ -367,7 +368,12 @@ Y_UNIT_TEST_SUITE(KqpKv) {
         CompareYson(Sprintf("[[[%du];[%du]]]", valueToReturn_1, valueToReturn_2), TString{res});
     }
 
-    Y_UNIT_TEST_TWIN(ReadRows_ExternalBlobs, NewPrecharge) {
+    template <typename FutureT>
+    auto GetValue(TKikimrRunner& kikimr, FutureT&& future) {
+        return kikimr.GetTestServer().GetRuntime()->WaitFuture(future);
+    }
+
+    Y_UNIT_TEST_TWIN(ReadRows_ExternalBlobs, UseExtBlobsPrecharge) {
         NKikimrConfig::TImmediateControlsConfig controls;
 
         if (NewPrecharge) {
@@ -379,11 +385,16 @@ Y_UNIT_TEST_SUITE(KqpKv) {
         auto settings = TKikimrSettings()
             .SetFeatureFlags(flags)
             .SetWithSampleTables(false)
-            .SetControls(controls);
+            .SetUseRealThreads(false)
+            .SetControls(controls)
+            .SetGrpcServerOptions(NYdbGrpc::TServerOptions()
+                .SetGRpcShutdownDeadline(TDuration::MilliSeconds(0))
+            );
+
         auto kikimr = TKikimrRunner{settings};
 
         auto db = kikimr.GetTableClient();
-        auto session = db.CreateSession().GetValueSync().GetSession();
+        auto session = GetValue(kikimr, db.CreateSession()).GetSession();
         const auto tableName = "/Root/TestTable";
         const auto keyColumnName_1 = "blob_id";
         const auto keyColumnName_2 = "chunk_num";
@@ -400,7 +411,7 @@ Y_UNIT_TEST_SUITE(KqpKv) {
         builder.SetPrimaryKeyColumns({keyColumnName_1, keyColumnName_2});
         builder.AddNullableColumn(dataColumnName, EPrimitiveType::String);
 
-        auto result = session.CreateTable(tableName, builder.Build()).GetValueSync();
+        auto result = GetValue(kikimr, session.CreateTable(tableName, builder.Build()));
         UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
 
         TString largeValue(1_MB, 'L');
@@ -417,9 +428,27 @@ Y_UNIT_TEST_SUITE(KqpKv) {
         }
         rows.EndList();
 
-        auto upsertResult = db.BulkUpsert(tableName, rows.Build()).GetValueSync();
+        auto upsertResult = GetValue(kikimr, db.BulkUpsert(tableName, rows.Build()));
         UNIT_ASSERT_C(upsertResult.IsSuccess(), upsertResult.GetIssues().ToString());
 
+        auto server = &kikimr.GetTestServer();
+
+        WaitForCompaction(server, tableName);
+
+        ui32 blobRequestCount = 0;
+
+        auto holder = server->GetRuntime()->AddObserver<TEvBlobStorage::TEvGet>([&blobRequestCount](TEvBlobStorage::TEvGet::TPtr& ev) {
+            const auto& msg = ev->Get();
+            for (size_t i = 0; i < msg->QuerySize; i++) {
+                const auto& id = msg->Queries[i].Id;
+
+                if (id.BlobSize() == 1_MB + 8 /** ext blob meta */) {
+                    blobRequestCount++;
+                    break;
+                }
+            }
+        });
+
         NYdb::TValueBuilder keys;
         keys.BeginList();
         for (int i = 0; i < 10; i++) {
@@ -431,11 +460,7 @@ Y_UNIT_TEST_SUITE(KqpKv) {
         }
         keys.EndList();
 
-        auto server = &kikimr.GetTestServer();
-
-        WaitForCompaction(server, tableName);
-
-        auto selectResult = db.ReadRows(tableName, keys.Build()).GetValueSync();
+        auto selectResult = GetValue(kikimr, db.ReadRows(tableName, keys.Build()));
 
         UNIT_ASSERT_C(selectResult.IsSuccess(), selectResult.GetIssues().ToString());
 
@@ -444,10 +469,83 @@ Y_UNIT_TEST_SUITE(KqpKv) {
 
         UNIT_ASSERT(parser.TryNextRow());
 
-        auto val = parser.GetValue(0);
-        TValueParser valParser(val);
-        TUuidValue v = valParser.GetUuid();
-        Cout << v.ToString() << Endl;
+        UNIT_ASSERT_VALUES_EQUAL(UseExtBlobsPrecharge ? 2 : 10, blobRequestCount);
+    }
+
+    Y_UNIT_TEST(ReadRows_TimeoutCancelsReads) {
+        auto settings = TKikimrSettings()
+            .SetWithSampleTables(false)
+            .SetUseRealThreads(false)
+            .SetGrpcServerOptions(NYdbGrpc::TServerOptions()
+                .SetGRpcShutdownDeadline(TDuration::MilliSeconds(0))
+            );
+
+        auto kikimr = TKikimrRunner{settings};
+
+        auto db = kikimr.GetTableClient();
+        auto session = GetValue(kikimr, db.CreateSession()).GetSession();
+        const auto tableName = "/Root/TestTable";
+        const auto keyColumnName = "blob_id";
+        const auto dataColumnName = "data";
+
+        TTableBuilder builder;
+        builder.AddNonNullableColumn(keyColumnName, EPrimitiveType::Int32);
+        builder.SetPrimaryKeyColumn(keyColumnName);
+        builder.AddNullableColumn(dataColumnName, EPrimitiveType::String);
+        builder.SetUniformPartitions(2);
+        TExplicitPartitions partitions;
+        partitions.AppendSplitPoints(TValueBuilder().BeginTuple().AddElement().OptionalInt32(5).EndTuple().Build());
+        builder.SetPartitionAtKeys(partitions);
+
+        auto result = GetValue(kikimr, session.CreateTable(tableName, builder.Build()));
+        UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+        NYdb::TValueBuilder rows;
+        rows.BeginList();
+        for (int i = 0; i < 10; i++) {
+            rows.AddListItem()
+                .BeginStruct()
+                    .AddMember(keyColumnName).Int32(i)
+                    .AddMember(dataColumnName).String("foo")
+                .EndStruct();
+        }
+        rows.EndList();
+
+        auto upsertResult = GetValue(kikimr, db.BulkUpsert(tableName, rows.Build()));
+        UNIT_ASSERT_C(upsertResult.IsSuccess(), upsertResult.GetIssues().ToString());
+
+        ui32 cancelCount = 0;
+        bool droppedOneResult = false;
+
+        kikimr.GetTestServer().GetRuntime()->SetObserverFunc([&](TAutoPtr<IEventHandle>& ev) {
+            if (auto *evRead = ev->CastAsLocal<TEvDataShard::TEvReadResult>()) {
+                if (droppedOneResult) {
+                    // Drop one of two results, so only one cancel would be sent to the DataShard
+                    return TTestActorRuntime::EEventAction::PROCESS;
+                }
+                droppedOneResult = true;
+                return TTestActorRuntime::EEventAction::DROP;
+            }
+            if (auto *evReadCancel = ev->CastAsLocal<TEvDataShard::TEvReadCancel>()) {
+                cancelCount++;
+            }
+            return TTestActorRuntime::EEventAction::PROCESS;
+        });
+
+        NYdb::TValueBuilder keys;
+        keys.BeginList();
+        for (int i = 0; i < 10; i++) {
+            keys.AddListItem()
+                .BeginStruct()
+                    .AddMember(keyColumnName).Int32(i)
+                .EndStruct();
+        }
+        keys.EndList();
+
+        auto selectResult = GetValue(kikimr, db.ReadRows(tableName, keys.Build()));
+
+        UNIT_ASSERT(!selectResult.IsSuccess());
+        UNIT_ASSERT_VALUES_EQUAL(cancelCount, 1);
     }
 
     TVector<::ReadRowsPgParam> readRowsPgParams
diff --git a/ydb/core/tx/datashard/datashard__read_iterator.cpp b/ydb/core/tx/datashard/datashard__read_iterator.cpp
@@ -3365,6 +3365,8 @@ void TDataShard::Handle(TEvDataShard::TEvReadCancel::TPtr& ev, const TActorConte
         state.Request->ReadSpan.EndError("Cancelled");
     }
     DeleteReadIterator(it);
+
+    LOG_WARN_S(ctx, NKikimrServices::TX_DATASHARD, TabletID() << " Cancelled read: " << readId);
 }
 
 void TDataShard::Handle(TEvDataShard::TEvReadScanStarted::TPtr& ev) {

Original file line number	Diff line number	Diff line change
`@@ -90,6 +90,7 @@ struct TKikimrSettings: public TTestFeatureFlagsHolder<TKikimrSettings> {`
`90`	`90`	`NMonitoring::TDynamicCounterPtr CountersRoot = MakeIntrusive<NMonitoring::TDynamicCounters>();`
`91`	`91`	`std::shared_ptr<NYql::NDq::IS3ActorsFactory> S3ActorsFactory = NYql::NDq::CreateDefaultS3ActorsFactory();`
`92`	`92`	`NKikimrConfig::TImmediateControlsConfig Controls;`
	`93`	`+ TMaybe<NYdbGrpc::TServerOptions> GrpcServerOptions;`
`93`	`94`
`94`	`95`	`TKikimrSettings()`
`95`	`96`	`{`
`@@ -128,6 +129,7 @@ struct TKikimrSettings: public TTestFeatureFlagsHolder<TKikimrSettings> {`
`128`	`129`	`return *this;`
`129`	`130`	`}`
`130`	`131`	`TKikimrSettings& SetControls(const NKikimrConfig::TImmediateControlsConfig& value) { Controls = value; return *this; }`
	`132`	`+ TKikimrSettings& SetGrpcServerOptions(const NYdbGrpc::TServerOptions& grpcServerOptions) { GrpcServerOptions = grpcServerOptions; return *this; };`
`131`	`133`	`};`
`132`	`134`
`133`	`135`	`class TKikimrRunner {`
Original file line number	Diff line number	Diff line change
`@@ -3365,6 +3365,8 @@ void TDataShard::Handle(TEvDataShard::TEvReadCancel::TPtr& ev, const TActorConte`
`3365`	`3365`	`state.Request->ReadSpan.EndError("Cancelled");`
`3366`	`3366`	`}`
`3367`	`3367`	`DeleteReadIterator(it);`
	`3368`	`+`
	`3369`	`+ LOG_WARN_S(ctx, NKikimrServices::TX_DATASHARD, TabletID() << " Cancelled read: " << readId);`
`3368`	`3370`	`}`
`3369`	`3371`
`3370`	`3372`	`void TDataShard::Handle(TEvDataShard::TEvReadScanStarted::TPtr& ev) {`