Draft ES-10826

Author: tlrx

server/src/main/java/org/elasticsearch/index/shard/EngineRef.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.shard;
+
+import org.elasticsearch.core.AbstractRefCounted;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.index.engine.Engine;
+
+public final class EngineRef extends AbstractRefCounted implements Releasable {
+
+    private final Releasable releasable;
+    private final Engine engine;
+
+    public EngineRef(Engine engine, Releasable releasable) {
+        this.engine = engine;
+        this.releasable = releasable;
+    }
+
+    @Nullable
+    public Engine getEngineOrNull() {
+        return engine;
+    }
+
+    @Override
+    protected void closeInternal() {
+        releasable.close();
+    }
+
+    @Override
+    public void close() {
+        decRef();
+    }
+}

server/src/main/java/org/elasticsearch/index/shard/EngineReferenceManager.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.shard;
+
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.index.engine.Engine;
+
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.Supplier;
+
+class EngineReferenceManager {
+
+    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(true); // fair
+    private final Releasable releaseExclusiveLock = () -> lock.writeLock().unlock();  // reuse this to avoid allocation for each op
+    private final Releasable releaseLock = () -> lock.readLock().unlock();            // reuse this to avoid allocation for each op
+
+    private volatile Engine current;
+
+    private Releasable acquireLock() {
+        lock.readLock().lock();
+        return Releasables.assertOnce(releaseLock);
+    }
+
+    public Releasable acquireExclusiveLock() {
+        lock.writeLock().lock();
+        return Releasables.assertOnce(releaseExclusiveLock);
+    }
+
+    public boolean isExclusiveLockHeldByCurrentThread() {
+        return lock.writeLock().isHeldByCurrentThread();
+    }
+
+    public EngineRef getEngineRef() {
+        try (var ignored = acquireLock()) {
+            return new EngineRef(this.current, acquireLock());
+        }
+    }
+
+    @Nullable
+    public Engine getAndSet(Supplier<Engine> supplier) {
+        assert supplier != null : "supplier cannot be null";
+        assert isExclusiveLockHeldByCurrentThread();
+        Engine previous = this.current;
+        this.current = supplier.get();
+        return previous;
+    }
+}

server/src/main/java/org/elasticsearch/index/shard/IndexShard.java

@@ -240,8 +240,8 @@ public class IndexShard extends AbstractIndexShardComponent implements IndicesCl
     // ensure happens-before relation between addRefreshListener() and postRecovery()
     private volatile SubscribableListener<Void> postRecoveryComplete;
     private volatile long pendingPrimaryTerm; // see JavaDocs for getPendingPrimaryTerm
-    private final Object engineMutex = new Object(); // lock ordering: engineMutex -> mutex
-    private final AtomicReference<Engine> currentEngineReference = new AtomicReference<>();
+
+    private final EngineReferenceManager engineReferenceManager = new EngineReferenceManager();
     final EngineFactory engineFactory;
 
     private final IndexingOperationListener indexingOperationListeners;
@@ -1613,10 +1613,10 @@ public Store.MetadataSnapshot snapshotStoreMetadata() throws IOException {
         Engine.IndexCommitRef indexCommit = null;
         store.incRef();
         try {
-            synchronized (engineMutex) {
+            try (var engineRef = engineReferenceManager.getEngineRef()) {
                 // if the engine is not running, we can access the store directly, but we need to make sure no one starts
                 // the engine on us. If the engine is running, we can get a snapshot via the deletion policy of the engine.
-                final Engine engine = getEngineOrNull();
+                final Engine engine = engineRef.getEngineOrNull();
                 if (engine != null) {
                     indexCommit = engine.acquireLastIndexCommit(false);
                 }
@@ -1776,14 +1776,14 @@ public CacheHelper getReaderCacheHelper() {
     }
 
     public void close(String reason, boolean flushEngine, Executor closeExecutor, ActionListener<Void> closeListener) throws IOException {
-        synchronized (engineMutex) {
+        try (var ignored = engineReferenceManager.acquireExclusiveLock()) {
             try {
                 synchronized (mutex) {
                     changeState(IndexShardState.CLOSED, reason);
                 }
                 checkAndCallWaitForEngineOrClosedShardListeners();
             } finally {
-                final Engine engine = this.currentEngineReference.getAndSet(null);
+                final Engine engine = engineReferenceManager.getAndSet(() -> null);
                 closeExecutor.execute(ActionRunnable.run(closeListener, new CheckedRunnable<>() {
                     @Override
                     public void run() throws Exception {
@@ -1857,7 +1857,7 @@ public void prepareForIndexRecovery() {
             throw new IndexShardNotRecoveringException(shardId, state);
         }
         recoveryState.setStage(RecoveryState.Stage.INDEX);
-        assert currentEngineReference.get() == null;
+        assert assertEngineReferenceIsNull("prepare for recovery, engine should be null");
     }
 
     /**
@@ -1936,8 +1936,8 @@ private void doLocalRecovery(
             // First, start a temporary engine, recover the local translog up to the given checkpoint, and then close the engine again.
             .<Void>newForked(l -> ActionListener.runWithResource(ActionListener.assertOnce(l), () -> () -> {
                 assert Thread.holdsLock(mutex) == false : "must not hold the mutex here";
-                synchronized (engineMutex) {
-                    IOUtils.close(currentEngineReference.getAndSet(null));
+                try (var ignored = engineReferenceManager.acquireExclusiveLock()) {
+                    IOUtils.close(engineReferenceManager.getAndSet(() -> null));
                 }
             }, (recoveryCompleteListener, ignoredRef) -> {
                 assert Thread.holdsLock(mutex) == false : "must not hold the mutex here";
@@ -2167,13 +2167,14 @@ private void innerOpenEngineAndTranslog(LongSupplier globalCheckpointSupplier) t
                 + recoveryState.getRecoverySource()
                 + "] but got "
                 + getRetentionLeases();
-        synchronized (engineMutex) {
-            assert currentEngineReference.get() == null : "engine is running";
+        try (var ignored = engineReferenceManager.acquireExclusiveLock()) {
+            assert assertEngineReferenceIsNull("engine is running");
             verifyNotClosed();
-            // we must create a new engine under mutex (see IndexShard#snapshotStoreMetadata).
+            // we must create a new engine under lock (see IndexShard#snapshotStoreMetadata).
             final Engine newEngine = createEngine(config);
             onNewEngine(newEngine);
-            currentEngineReference.set(newEngine);
+            var previous = engineReferenceManager.getAndSet(() -> newEngine);
+            assert previous == null;
             // We set active because we are now writing operations to the engine; this way,
             // we can flush if we go idle after some time and become inactive.
             active.set(true);
@@ -2241,7 +2242,7 @@ private boolean assertLastestCommitUserData() throws IOException {
     }
 
     private void onNewEngine(Engine newEngine) {
-        assert Thread.holdsLock(engineMutex);
+        assert engineReferenceManager.isExclusiveLockHeldByCurrentThread();
         refreshListeners.setCurrentRefreshLocationSupplier(newEngine::getTranslogLastWriteLocation);
         refreshListeners.setCurrentProcessedCheckpointSupplier(newEngine::getProcessedLocalCheckpoint);
         refreshListeners.setMaxIssuedSeqNoSupplier(newEngine::getMaxSeqNo);
@@ -2252,9 +2253,9 @@ private void onNewEngine(Engine newEngine) {
      */
     public void performRecoveryRestart() throws IOException {
         assert Thread.holdsLock(mutex) == false : "restart recovery under mutex";
-        synchronized (engineMutex) {
+        try (var ignored = engineReferenceManager.acquireExclusiveLock()) {
             assert refreshListeners.pendingCount() == 0 : "we can't restart with pending listeners";
-            IOUtils.close(currentEngineReference.getAndSet(null));
+            IOUtils.close(engineReferenceManager.getAndSet(() -> null));
             resetRecoveryStage();
         }
     }
@@ -2264,7 +2265,7 @@ public void performRecoveryRestart() throws IOException {
      */
     public void resetRecoveryStage() {
         assert routingEntry().recoverySource().getType() == RecoverySource.Type.PEER : "not a peer recovery [" + routingEntry() + "]";
-        assert currentEngineReference.get() == null;
+        assert assertEngineReferenceIsNull("reset recovery stage, engine should be null");
         if (state != IndexShardState.RECOVERING) {
             throw new IndexShardNotRecoveringException(shardId, state);
         }
@@ -3297,8 +3298,11 @@ Engine getEngine() {
      * NOTE: returns null if engine is not yet started (e.g. recovery phase 1, copying over index files, is still running), or if engine is
      * closed.
      */
+    @Deprecated
     public Engine getEngineOrNull() {
-        return this.currentEngineReference.get();
+        try (var engineRef = engineReferenceManager.getEngineRef()) {
+            return engineRef.getEngineOrNull();
+        }
     }
 
     public void startRecovery(
@@ -4312,11 +4316,11 @@ public void resetEngine() {
         assert Thread.holdsLock(mutex) == false : "resetting engine under mutex";
         assert waitForEngineOrClosedShardListeners.isDone();
         try {
-            synchronized (engineMutex) {
+            try (var ignored = engineReferenceManager.acquireExclusiveLock()) {
                 verifyNotClosed();
                 getEngine().prepareForEngineReset();
                 var newEngine = createEngine(newEngineConfig(replicationTracker));
-                IOUtils.close(currentEngineReference.getAndSet(newEngine));
+                IOUtils.close(engineReferenceManager.getAndSet(() -> newEngine));
                 onNewEngine(newEngine);
             }
             onSettingsChanged();
@@ -4342,7 +4346,7 @@ assert getActiveOperationsCount() == OPERATIONS_BLOCKED
         SetOnce<Engine> newEngineReference = new SetOnce<>();
         final long globalCheckpoint = getLastKnownGlobalCheckpoint();
         assert globalCheckpoint == getLastSyncedGlobalCheckpoint();
-        synchronized (engineMutex) {
+        try (var ignored = engineReferenceManager.acquireExclusiveLock()) {
             verifyNotClosed();
             // we must create both new read-only engine and new read-write engine under engineMutex to ensure snapshotStoreMetadata,
             // acquireXXXCommit and close works.
@@ -4357,7 +4361,7 @@ assert getActiveOperationsCount() == OPERATIONS_BLOCKED
             ) {
                 @Override
                 public IndexCommitRef acquireLastIndexCommit(boolean flushFirst) {
-                    synchronized (engineMutex) {
+                    try (var ignored = engineReferenceManager.getEngineRef()) {
                         if (newEngineReference.get() == null) {
                             throw new AlreadyClosedException("engine was closed");
                         }
@@ -4368,7 +4372,7 @@ public IndexCommitRef acquireLastIndexCommit(boolean flushFirst) {
 
                 @Override
                 public IndexCommitRef acquireSafeIndexCommit() {
-                    synchronized (engineMutex) {
+                    try (var ignored = engineReferenceManager.getEngineRef()) {
                         if (newEngineReference.get() == null) {
                             throw new AlreadyClosedException("engine was closed");
                         }
@@ -4379,17 +4383,17 @@ public IndexCommitRef acquireSafeIndexCommit() {
                 @Override
                 public void close() throws IOException {
                     Engine newEngine;
-                    synchronized (engineMutex) {
+                    try (var engineRef = engineReferenceManager.getEngineRef()) {
                         newEngine = newEngineReference.get();
-                        if (newEngine == currentEngineReference.get()) {
+                        if (newEngine == engineRef.getEngineOrNull()) {
                             // we successfully installed the new engine so do not close it.
                             newEngine = null;
                         }
                     }
                     IOUtils.close(super::close, newEngine);
                 }
             };
-            IOUtils.close(currentEngineReference.getAndSet(readOnlyEngine));
+            IOUtils.close(engineReferenceManager.getAndSet(() -> readOnlyEngine));
             newEngineReference.set(engineFactory.newReadWriteEngine(newEngineConfig(replicationTracker)));
             onNewEngine(newEngineReference.get());
         }
@@ -4403,9 +4407,9 @@ public void close() throws IOException {
         );
         newEngineReference.get().recoverFromTranslog(translogRunner, globalCheckpoint);
         newEngineReference.get().refresh("reset_engine");
-        synchronized (engineMutex) {
+        try (var ignored = engineReferenceManager.acquireExclusiveLock()) {
             verifyNotClosed();
-            IOUtils.close(currentEngineReference.getAndSet(newEngineReference.get()));
+            IOUtils.close(engineReferenceManager.getAndSet(() -> newEngineReference.get()));
             // We set active because we are now writing operations to the engine; this way,
             // if we go idle after some time and become inactive, we still give sync'd flush a chance to run.
             active.set(true);
@@ -4516,4 +4520,11 @@ public void ensureMutable(ActionListener<Void> listener) {
             l.onResponse(null);
         }));
     }
+
+    private boolean assertEngineReferenceIsNull(String message)  {
+        try (var engineRef = engineReferenceManager.getEngineRef()) {
+            assert engineRef.getEngineOrNull() == null : message;
+            return true;
+        }
+    }
 }