[CI] Fix IndexShardTests.testReentrantEngineReadLockAcquisitionInRefreshListener (#126685)

I suspect the test resets/closes the reference manager
between the refresh and the retrieval of the segment
generation after the refresh.

By executing segmentGenerationAfterRefresh while
holding the engine reset lock we make sure there
are no concurrent engine resets meanwhile.

In the future, we should also ensure that
IndexShard.refresh() uses withEngine.

Closes #126628

Author: tlrx

muted-tests.yml

@@ -411,9 +411,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));