Tighten up ref counting on recovery target (elastic#113767)

Today we throw an exception in some places if the `RecoveryTarget` is
used after all refs are released, but really (a) it's a bug to hit this
exception and (b) the consequences of such a bug are mostly not
especially severe.

This commit replaces the exception with an assertion, and acquires refs
properly in the few places that might not hold a ref already today.

Author: DaveCTurner

server/src/main/java/org/elasticsearch/indices/recovery/MultiFileWriter.java

@@ -42,7 +42,6 @@ public class MultiFileWriter extends AbstractRefCounted implements Releasable {
     private final RecoveryState.Index indexState;
     private final String tempFilePrefix;
     private final Logger logger;
-    private final Runnable ensureOpen;
     private final boolean verifyOutput;
 
     private final AtomicBoolean closed = new AtomicBoolean(false);
@@ -52,23 +51,15 @@ public class MultiFileWriter extends AbstractRefCounted implements Releasable {
 
     final Map<String, String> tempFileNames = ConcurrentCollections.newConcurrentMap();
 
-    public MultiFileWriter(Store store, RecoveryState.Index indexState, String tempFilePrefix, Logger logger, Runnable ensureOpen) {
-        this(store, indexState, tempFilePrefix, logger, ensureOpen, true);
+    public MultiFileWriter(Store store, RecoveryState.Index indexState, String tempFilePrefix, Logger logger) {
+        this(store, indexState, tempFilePrefix, logger, true);
     }
 
-    public MultiFileWriter(
-        Store store,
-        RecoveryState.Index indexState,
-        String tempFilePrefix,
-        Logger logger,
-        Runnable ensureOpen,
-        boolean verifyOutput
-    ) {
+    public MultiFileWriter(Store store, RecoveryState.Index indexState, String tempFilePrefix, Logger logger, boolean verifyOutput) {
         this.store = store;
         this.indexState = indexState;
         this.tempFilePrefix = tempFilePrefix;
         this.logger = logger;
-        this.ensureOpen = ensureOpen;
         this.verifyOutput = verifyOutput;
     }
 
@@ -85,7 +76,7 @@ public void writeFileChunk(StoreFileMetadata fileMetadata, long position, Releas
     }
 
     public void writeFile(StoreFileMetadata fileMetadata, long readSnapshotFileBufferSize, InputStream stream) throws IOException {
-        ensureOpen.run();
+        assert hasReferences();
         assert Transports.assertNotTransportThread("multi_file_writer");
 
         String fileName = fileMetadata.name();
@@ -146,13 +137,13 @@ String getTempNameForFile(String origFile) {
     }
 
     public IndexOutput getOpenIndexOutput(String key) {
-        ensureOpen.run();
+        assert hasReferences();
         return openIndexOutputs.get(key);
     }
 
     /** remove and {@link IndexOutput} for a given file. It is the caller's responsibility to close it */
     public IndexOutput removeOpenIndexOutputs(String name) {
-        ensureOpen.run();
+        assert hasReferences();
         return openIndexOutputs.remove(name);
     }
 
@@ -164,7 +155,7 @@ public IndexOutput removeOpenIndexOutputs(String name) {
      * at a later stage
      */
     public IndexOutput openAndPutIndexOutput(String fileName, StoreFileMetadata metadata, Store store) throws IOException {
-        ensureOpen.run();
+        assert hasReferences();
         String tempFileName = getTempNameForFile(fileName);
         if (tempFileNames.containsKey(tempFileName)) {
             throw new IllegalStateException("output for file [" + fileName + "] has already been created");
@@ -178,6 +169,7 @@ public IndexOutput openAndPutIndexOutput(String fileName, StoreFileMetadata meta
 
     private void innerWriteFileChunk(StoreFileMetadata fileMetadata, long position, BytesReference content, boolean lastChunk)
         throws IOException {
+        assert hasReferences();
         final String name = fileMetadata.name();
         IndexOutput indexOutput;
         if (position == 0) {
@@ -242,7 +234,7 @@ protected void closeInternal() {
 
     /** renames all temporary files to their true name, potentially overriding existing files */
     public void renameAllTempFiles() throws IOException {
-        ensureOpen.run();
+        assert hasReferences();
         store.renameTempFilesSafe(tempFileNames);
     }
 

server/src/main/java/org/elasticsearch/indices/recovery/PeerRecoveryTargetService.java

@@ -401,8 +401,13 @@ record StartRecoveryRequestToSend(StartRecoveryRequest startRecoveryRequest, Str
                 .andThenApply(startingSeqNo -> {
                     assert startingSeqNo == UNASSIGNED_SEQ_NO || recoveryTarget.state().getStage() == RecoveryState.Stage.TRANSLOG
                         : "unexpected recovery stage [" + recoveryTarget.state().getStage() + "] starting seqno [ " + startingSeqNo + "]";
-                    final var startRequest = getStartRecoveryRequest(logger, clusterService.localNode(), recoveryTarget, startingSeqNo);
-                    return new StartRecoveryRequestToSend(startRequest, PeerRecoverySourceService.Actions.START_RECOVERY, startRequest);
+                    try {
+                        recoveryTarget.incRef();
+                        final var startRequest = getStartRecoveryRequest(logger, clusterService.localNode(), recoveryTarget, startingSeqNo);
+                        return new StartRecoveryRequestToSend(startRequest, PeerRecoverySourceService.Actions.START_RECOVERY, startRequest);
+                    } finally {
+                        recoveryTarget.decRef();
+                    }
                 })
                 // finally send the start-recovery request
                 .addListener(toSendListener);
@@ -572,7 +577,14 @@ protected void handleRequest(
 
         @Override
         protected CheckedFunction<Void, TransportResponse, Exception> responseMapping(RecoveryTarget recoveryTarget) {
-            return v -> new RecoveryTranslogOperationsResponse(recoveryTarget.indexShard().getLocalCheckpoint());
+            return v -> {
+                try {
+                    recoveryTarget.incRef();
+                    return new RecoveryTranslogOperationsResponse(recoveryTarget.indexShard().getLocalCheckpoint());
+                } finally {
+                    recoveryTarget.decRef();
+                }
+            };
         }
 
         private void performTranslogOps(

server/src/main/java/org/elasticsearch/indices/recovery/RecoveryTarget.java

@@ -13,7 +13,6 @@
 import org.apache.lucene.index.CorruptIndexException;
 import org.apache.lucene.index.IndexFormatTooNewException;
 import org.apache.lucene.index.IndexFormatTooOldException;
-import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.admin.indices.flush.FlushRequest;
@@ -142,7 +141,7 @@ private void recreateMultiFileWriter() {
 
     private MultiFileWriter createMultiFileWriter() {
         final String tempFilePrefix = RECOVERY_PREFIX + UUIDs.randomBase64UUID() + ".";
-        return new MultiFileWriter(indexShard.store(), indexShard.recoveryState().getIndex(), tempFilePrefix, logger, this::ensureRefCount);
+        return new MultiFileWriter(indexShard.store(), indexShard.recoveryState().getIndex(), tempFilePrefix, logger);
     }
 
     /**
@@ -179,7 +178,7 @@ public ShardId shardId() {
     }
 
     public IndexShard indexShard() {
-        ensureRefCount();
+        assert hasReferences();
         return indexShard;
     }
 
@@ -232,7 +231,7 @@ public Releasable disableRecoveryMonitor() {
     }
 
     public Store store() {
-        ensureRefCount();
+        assert hasReferences();
         return store;
     }
 
@@ -361,14 +360,6 @@ public String toString() {
         return shardId + " [" + recoveryId + "]";
     }
 
-    private void ensureRefCount() {
-        if (refCount() <= 0) {
-            throw new ElasticsearchException(
-                "RecoveryStatus is used but it's refcount is 0. Probably a mismatch between incRef/decRef " + "calls"
-            );
-        }
-    }
-
     /*** Implementation of {@link RecoveryTargetHandler } */
 
     @Override

server/src/test/java/org/elasticsearch/indices/recovery/MultiFileWriterTests.java

@@ -90,7 +90,7 @@ public void testWritesFileWithIncorrectChecksumWithoutVerification() throws IOEx
     }
 
     private MultiFileWriter createMultiFileWriter(boolean verifyOutput) {
-        return new MultiFileWriter(store, mock(RecoveryState.Index.class), "temp_", logger, mock(Runnable.class), verifyOutput);
+        return new MultiFileWriter(store, mock(RecoveryState.Index.class), "temp_", logger, verifyOutput);
     }
 
     private record FileAndMetadata(byte[] bytes, StoreFileMetadata metadata) {}

server/src/test/java/org/elasticsearch/indices/recovery/RecoverySourceHandlerTests.java

@@ -191,7 +191,7 @@ public void testSendFiles() throws Throwable {
             metas.add(md);
         }
         Store targetStore = newStore(createTempDir());
-        MultiFileWriter multiFileWriter = new MultiFileWriter(targetStore, mock(RecoveryState.Index.class), "", logger, () -> {});
+        MultiFileWriter multiFileWriter = new MultiFileWriter(targetStore, mock(RecoveryState.Index.class), "", logger);
         RecoveryTargetHandler target = new TestRecoveryTargetHandler() {
             @Override
             public void writeFileChunk(
@@ -578,7 +578,7 @@ public void testHandleCorruptedIndexOnSendSendFiles() throws Throwable {
             )
         );
         Store targetStore = newStore(createTempDir(), false);
-        MultiFileWriter multiFileWriter = new MultiFileWriter(targetStore, mock(RecoveryState.Index.class), "", logger, () -> {});
+        MultiFileWriter multiFileWriter = new MultiFileWriter(targetStore, mock(RecoveryState.Index.class), "", logger);
         RecoveryTargetHandler target = new TestRecoveryTargetHandler() {
             @Override
             public void writeFileChunk(

server/src/test/java/org/elasticsearch/indices/recovery/RecoveryStatusTests.java

@@ -36,8 +36,7 @@ public void testRenameTempFiles() throws IOException {
             indexShard.store(),
             indexShard.recoveryState().getIndex(),
             "recovery.test.",
-            logger,
-            () -> {}
+            logger
         );
         try (
             IndexOutput indexOutput = multiFileWriter.openAndPutIndexOutput(

server/src/test/java/org/elasticsearch/recovery/RecoveriesCollectionTests.java

@@ -8,8 +8,8 @@
  */
 package org.elasticsearch.recovery;
 
-import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.core.Assertions;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.replication.ESIndexLevelReplicationTestCase;
 import org.elasticsearch.index.shard.IndexShard;
@@ -139,8 +139,10 @@ public void testResetRecovery() throws Exception {
             assertEquals(referencesToStore, resetRecovery.store().refCount());
             assertEquals(currentAsTarget, shard.recoveryStats().currentAsTarget());
             assertEquals(recoveryTarget.refCount(), 0);
-            expectThrows(ElasticsearchException.class, () -> recoveryTarget.store());
-            expectThrows(ElasticsearchException.class, () -> recoveryTarget.indexShard());
+            if (Assertions.ENABLED) {
+                expectThrows(AssertionError.class, recoveryTarget::store);
+                expectThrows(AssertionError.class, recoveryTarget::indexShard);
+            }
             String resetTempFileName = resetRecovery.getTempNameForFile("foobar");
             assertNotEquals(tempFileName, resetTempFileName);
             assertEquals(currentAsTarget, shard.recoveryStats().currentAsTarget());

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/repository/CcrRepository.java

@@ -740,7 +740,7 @@ protected void restoreFiles(List<FileInfo> filesToRecover, Store store, ActionLi
                 mds
             ) {
 
-                final MultiFileWriter multiFileWriter = new MultiFileWriter(store, recoveryState.getIndex(), "", logger, () -> {});
+                final MultiFileWriter multiFileWriter = new MultiFileWriter(store, recoveryState.getIndex(), "", logger);
                 long offset = 0;
 
                 @Override