[CI] Fix IndexShardTests.testReentrantEngineReadLockAcquisitionInRefreshListener

muted-tests.yml

@@ -414,9 +414,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.action.ForkIT
   method: testWithStatsSimple
   issue: https://github.com/elastic/elasticsearch/issues/126607
-- class: org.elasticsearch.index.shard.IndexShardTests
-  method: testReentrantEngineReadLockAcquisitionInRefreshListener
-  issue: https://github.com/elastic/elasticsearch/issues/126628
 - class: org.elasticsearch.xpack.esql.action.EsqlActionIT
   method: testQueryOnEmptyDataIndex
   issue: https://github.com/elastic/elasticsearch/issues/126580

server/src/main/java/org/elasticsearch/index/engine/InternalEngine.java

@@ -2063,6 +2063,7 @@ protected final RefreshResult refresh(String source, SearcherScope scope, boolea
                         engineReadLock.lock();
                         try {
                             referenceManager.maybeRefreshBlocking();
+                            segmentGeneration = segmentGenerationAfterRefresh(referenceManager, generationBeforeRefresh);
                             refreshed = true;
                         } finally {
                             engineReadLock.unlock();
@@ -2071,20 +2072,14 @@ protected final RefreshResult refresh(String source, SearcherScope scope, boolea
                         if (engineReadLock.tryLock()) {
                             try {
                                 refreshed = referenceManager.maybeRefresh();
+                                if (refreshed) {
+                                    segmentGeneration = segmentGenerationAfterRefresh(referenceManager, generationBeforeRefresh);
+                                }
                             } finally {
                                 engineReadLock.unlock();
                             }
                         }
                     }
-                    if (refreshed) {
-                        final ElasticsearchDirectoryReader current = referenceManager.acquire();
-                        try {
-                            // Just use the generation from the reader when https://github.com/apache/lucene/pull/12177 is included.
-                            segmentGeneration = Math.max(current.getIndexCommit().getGeneration(), generationBeforeRefresh);
-                        } finally {
-                            referenceManager.release(current);
-                        }
-                    }
                 } finally {
                     store.decRef();
                 }
@@ -2120,6 +2115,21 @@ protected final RefreshResult refresh(String source, SearcherScope scope, boolea
         return new RefreshResult(refreshed, primaryTerm, segmentGeneration);
     }
 
+    private long segmentGenerationAfterRefresh(
+        ReferenceManager<ElasticsearchDirectoryReader> referenceManager,
+        long generationBeforeRefresh
+    ) throws IOException {
+        assert store.hasReferences();
+        assert engineConfig.getEngineResetLock().isReadLockedByCurrentThread() : "prevent concurrent engine resets";
+        final ElasticsearchDirectoryReader current = referenceManager.acquire();
+        try {
+            // Just use the generation from the reader when https://github.com/apache/lucene/pull/12177 is included.
+            return Math.max(current.getIndexCommit().getGeneration(), generationBeforeRefresh);
+        } finally {
+            referenceManager.release(current);
+        }
+    }
+
     @Override
     public void writeIndexingBuffer() throws IOException {
         final long reclaimableVersionMapBytes = versionMap.reclaimableRefreshRamBytes();

server/src/test/java/org/elasticsearch/index/shard/IndexShardTests.java

@@ -5274,7 +5274,7 @@ public void beforeRefresh() throws IOException {
                         refreshStarted.countDown();
                         safeAwait(getFromTranslogStarted);
 
-                        // A this stage, getThread is blocked on the refresh held by the current thread
+                        // A this stage, getThread is blocked on the refresh lock held by the current thread
                         assertBusy(() -> assertThat(engineResetLock.getReadLockCount(), greaterThanOrEqualTo(2)));
                         assertThat(getFromTranslogResult.isDone(), equalTo(false));