Fix infinite retries

1 files changed, 15 insertions, 3 deletions

diff --git a/ydb/core/kqp/runtime/kqp_read_actor.cpp b/ydb/core/kqp/runtime/kqp_read_actor.cpp
index 472c6f25c54..47c494edc3e 100644
--- a/ydb/core/kqp/runtime/kqp_read_actor.cpp
+++ b/ydb/core/kqp/runtime/kqp_read_actor.cpp
@@ -22,7 +22,6 @@ static constexpr ui64 EVREAD_MAX_ROWS = 32767;
 static constexpr ui64 EVREAD_MAX_BYTES = 200_MB;
 
 static constexpr ui64 MAX_SHARD_RETRIES = 5;
-//static constexpr ui64 MAX_TOTAL_SHARD_RETRIES = 20;
 static constexpr ui64 MAX_SHARD_RESOLVES = 3;
 
 bool IsDebugLogEnabled(const NActors::TActorSystem* actorSystem, NActors::NLog::EComponent component) {
@@ -56,6 +55,8 @@ public:
         size_t ResolveAttempt = 0;
         size_t RetryAttempt = 0;
 
+        bool NeedResolve = false;
+
         TShardState(ui64 tabletId)
             : TabletId(tabletId)
         {
@@ -251,7 +252,6 @@ public:
     }
 
     void Bootstrap() {
-        //TODO: resolve if hint is not set
         THolder<TShardState> stateHolder = MakeHolder<TShardState>(Settings.GetShardIdHint());
         PendingShards.PushBack(stateHolder.Get());
         auto& state = *stateHolder.Release();
@@ -272,7 +272,12 @@ public:
             }
         }
 
-        StartTableScan();
+        if (!Settings.HasShardIdHint()) {
+            state.NeedResolve = true;
+            ResolveShard(&state);
+        } else {
+            StartTableScan();
+        }
         Become(&TKqpReadActor::ReadyState);
     }
 
@@ -389,6 +394,13 @@ public:
             return;
         }
 
+        if (keyDesc->GetPartitions().size() == 1 && !state->NeedResolve) {
+            // we re-resolved the same shard
+            RuntimeError(TStringBuilder() << "too many retries for shard " << state->TabletId, NDqProto::StatusIds::StatusIds::INTERNAL_ERROR);
+            PendingShards.PushBack(state.Release());
+            return;
+        }
+
         if (keyDesc->GetPartitions().empty()) {
             TString error = TStringBuilder() << "No partitions to read from '" << Settings.GetTable().GetTablePath() << "'";
             CA_LOG_E(error);