Fix request processing scheduling

Author: idegtiarenko

muted-tests.yml

@@ -417,9 +417,6 @@ tests:
 - class: org.elasticsearch.packaging.test.DockerTests
   method: test026InstallBundledRepositoryPluginsViaConfigFile
   issue: https://github.com/elastic/elasticsearch/issues/127158
-- class: org.elasticsearch.xpack.esql.plugin.DataNodeRequestSenderTests
-  method: testRetryOnlyMovedShards
-  issue: https://github.com/elastic/elasticsearch/issues/127168
 - class: org.elasticsearch.xpack.esql.plugin.DataNodeRequestSenderIT
   method: testSearchWhileRelocating
   issue: https://github.com/elastic/elasticsearch/issues/127188

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestSender.java

@@ -265,9 +265,8 @@ void onAfter(DriverCompletionInfo info) {
                     concurrentRequests.release();
                 }
 
-                if (pendingRetries.isEmpty() == false && remainingUnavailableShardResolutionAttempts.decrementAndGet() >= 0) {
+                if (pendingRetries.isEmpty() == false) {
                     try {
-                        sendingLock.lock();
                         var resolutions = resolveShards(pendingRetries);
                         for (var entry : resolutions.entrySet()) {
                             targetShards.shards.get(entry.getKey()).remainingNodes.addAll(entry.getValue());
@@ -292,8 +291,8 @@ public void onResponse(DataNodeComputeResponse response) {
                 for (var entry : response.shardLevelFailures().entrySet()) {
                     final ShardId shardId = entry.getKey();
                     trackShardLevelFailure(shardId, false, entry.getValue());
-                    pendingShardIds.add(shardId);
                     maybeScheduleRetry(shardId, false, entry.getValue());
+                    pendingShardIds.add(shardId);
                 }
                 onAfter(response.completionInfo());
             }
@@ -302,8 +301,8 @@ public void onResponse(DataNodeComputeResponse response) {
             public void onFailure(Exception e, boolean receivedData) {
                 for (ShardId shardId : request.shardIds) {
                     trackShardLevelFailure(shardId, receivedData, e);
-                    pendingShardIds.add(shardId);
                     maybeScheduleRetry(shardId, receivedData, e);
+                    pendingShardIds.add(shardId);
                 }
                 onAfter(DriverCompletionInfo.EMPTY);
             }
@@ -322,6 +321,9 @@ private void maybeScheduleRetry(ShardId shardId, boolean receivedData, Exception
                 if (receivedData == false
                     && targetShards.getShard(shardId).remainingNodes.isEmpty()
                     && unwrapFailure(shardId, e) instanceof NoShardAvailableActionException) {
+                    if (pendingRetries.isEmpty() && remainingUnavailableShardResolutionAttempts.decrementAndGet() >= 0) {
+                        sendingLock.lock();
+                    }
                     pendingRetries.add(shardId);
                 }
             }