Fix ydb read rows timeout after ds tablets restart (#17925)

Author: robdrynkin

ydb/core/grpc_services/rpc_read_rows.cpp

@@ -732,6 +732,14 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
         SendResult(status, errorMsg, issues);
     }
 
+    void Handle(TEvents::TEvUndelivered::TPtr&) {
+        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) {
+        return ReplyWithError(Ydb::StatusIds::UNAVAILABLE, TStringBuilder() << "Failed to connect to shard " << ev->Get()->TabletId);
+    }
+
     void PassAway() override {
         Send(PipeCache, new TEvPipeCache::TEvUnlink(0));
         if (TimeoutTimerActorId) {
@@ -748,6 +756,9 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
             hFunc(TEvTxProxySchemeCache::TEvResolveKeySetResult, Handle);
             hFunc(TEvDataShard::TEvReadResult, Handle);
 
+            hFunc(TEvents::TEvUndelivered, Handle);
+            hFunc(TEvPipeCache::TEvDeliveryProblem, Handle);
+
             hFunc(TEvents::TEvWakeup, HandleTimeout);
         }
     }

ydb/services/ydb/ut/ya.make

@@ -25,6 +25,7 @@ SRCS(
     ydb_olapstore_ut.cpp
     ydb_monitoring_ut.cpp
     ydb_query_ut.cpp
+    ydb_read_rows_ut.cpp
     ydb_ldap_login_ut.cpp
     ydb_login_ut.cpp
     ydb_object_storage_ut.cpp

ydb/services/ydb/ydb_read_rows_ut.cpp

@@ -0,0 +1,99 @@
+#include <ydb/core/grpc_services/base/base.h>
+
+#include <ydb/core/tx/datashard/defs.h>
+#include <ydb/core/tx/datashard/ut_common/datashard_ut_common.h>
+#include <ydb/core/tx/datashard/datashard_ut_common_kqp.h>
+
+#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/proto/accessor.h>
+
+#include <library/cpp/testing/unittest/registar.h>
+
+namespace NKikimr {
+
+using namespace NKikimr::NDataShard::NKqpHelpers;
+using namespace NSchemeShard;
+using namespace Tests;
+
+namespace {
+
+using TEvReadRowsRequest = NGRpcService::TGrpcRequestNoOperationCall<Ydb::Table::ReadRowsRequest, Ydb::Table::ReadRowsResponse>;
+
+using TRows = TVector<std::pair<TSerializedCellVec, TString>>;
+using TRowTypes = TVector<std::pair<TString, Ydb::Type>>;
+
+
+Ydb::Table::ReadRowsRequest MakeReadRowsRequest(const TString& tablePath, const TVector<ui32>& keys) {
+    Ydb::Table::ReadRowsRequest request;
+    request.set_path(tablePath);
+
+    NYdb::TValueBuilder keysBuilder;
+    keysBuilder.BeginList();
+    for (ui32 key : keys) {
+        keysBuilder.AddListItem().BeginStruct().AddMember("key").Uint32(key).EndStruct();
+    }
+    keysBuilder.EndList();
+
+    auto keysValuesCpp = keysBuilder.Build();
+    auto keysTypeCpp = keysValuesCpp.GetType();
+    request.mutable_keys()->mutable_type()->CopyFrom(NYdb::TProtoAccessor::GetProto(keysTypeCpp));
+    request.mutable_keys()->mutable_value()->CopyFrom(NYdb::TProtoAccessor::GetProto(keysValuesCpp));
+
+    return request;
+}
+
+} // namespace
+
+Y_UNIT_TEST_SUITE(ReadRows) {
+
+    Y_UNIT_TEST(KillTabletDuringRead) {
+        // Init cluster
+        TPortManager pm;
+        TServerSettings serverSettings(pm.GetPort(2134));
+        serverSettings.SetDomainName("Root")
+            .SetUseRealThreads(false);
+
+        Tests::TServer::TPtr server = new TServer(serverSettings);
+        auto &runtime = *server->GetRuntime();
+        auto sender = runtime.AllocateEdgeActor();
+
+        InitRoot(server, sender);
+
+        // Create table
+        CreateShardedTable(server, sender, "/Root", "table-1", 1, false);
+        ExecSQL(server, sender, "UPSERT INTO `/Root/table-1` (key, value) VALUES (1, 100), (3, 300), (5, 500);");
+
+        // Check normal behavior
+        {
+            Ydb::Table::ReadRowsRequest request = MakeReadRowsRequest("/Root/table-1", {1, 5});
+            auto readRowsFuture = NRpcService::DoLocalRpc<TEvReadRowsRequest>(
+                std::move(request), "/Root", "", runtime.GetActorSystem(0));
+            auto res = runtime.WaitFuture(readRowsFuture, TDuration::Seconds(10));
+            UNIT_ASSERT_VALUES_EQUAL(res.status(), ::Ydb::StatusIds::SUCCESS);
+        }
+
+        // Get tablet id of the only table shard
+        auto tablets = GetTableShards(server, sender, "/Root/table-1");
+        UNIT_ASSERT(tablets.size() == 1);
+        ui64 tabletId = tablets.at(0);
+
+        // Reboot tablet during read
+        auto dsReadResultOberver = runtime.AddObserver<TEvDataShard::TEvReadResult>([&](TEvDataShard::TEvReadResult::TPtr& ev) {
+            Cerr << "Stoping tablet id: " << tabletId;
+            RebootTablet(runtime, tabletId, sender);
+            ev.Reset();
+        });
+
+        // Check read with tablet reboot
+        {
+            Ydb::Table::ReadRowsRequest request = MakeReadRowsRequest("/Root/table-1", {1, 5});
+            auto readRowsFuture = NRpcService::DoLocalRpc<TEvReadRowsRequest>(
+                std::move(request), "/Root", "", runtime.GetActorSystem(0));
+            auto res = runtime.WaitFuture(readRowsFuture, TDuration::Seconds(10));
+            UNIT_ASSERT_VALUES_EQUAL(res.status(), ::Ydb::StatusIds::UNAVAILABLE);
+        }
+
+    }
+
+}
+
+} // namespace NKikimr