Introduce primary context (#25122)

* Introduce primary context

The target of a primary relocation is not aware of the state of the
replication group. In particular, it is not tracking in-sync and
initializing shards and their checkpoints. This means that after the
target shard is started, its knowledge of the replication group could
differ from that of the relocation source. In particular, this differing
view can lead to it computing a global checkpoint that moves backwards
after it becomes aware of the state of the entire replication
group. This commit addresses this issue by transferring a primary
context during relocation handoff.

* Fix test

* Add assertion messages

* Javadocs

* Barrier between marking a shard in sync and relocating

* Fix misplaced call

* Paranoia

* Better latch countdown

* Catch any exception

* Fix comment

* Fix wait for cluster state relocation test

* Update knowledge via upate local checkpoint API

* toString

* Visibility

* Refactor permit

* Push down

* Imports

* Docs

* Fix compilation

* Remove assertion

* Fix compilation

* Remove context wrapper

* Move PrimaryContext to new package

* Piping for cluster state version

This commit adds piping for the cluster state version to the global
checkpoint tracker. We do not use it yet.

* Remove unused import

* Implement versioning in tracker

* Fix test

* Unneeded public

* Imports

* Promote on our own

* Add tests

* Import

* Newline

* Update comment

* Serialization

* Assertion message

* Update stale comment

* Remove newline

* Less verbose

* Remove redundant assertion

* Tracking -> in-sync

* Assertions

* Just say no

Friends do not let friends block the cluster state update thread on
network operations.

* Extra newline

* Add allocation ID to assertion

* Rename method

* Another rename

* Introduce sealing

* Sealing tests

* One more assertion

* Fix imports

* Safer sealing

* Remove check

* Remove another sealed check

Author: jasontedor

core/src/main/java/org/elasticsearch/common/collect/LongTuple.java

@@ -0,0 +1,66 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.common.collect;
+
+public class LongTuple<T> {
+
+    public static <T> LongTuple<T> tuple(final T v1, final long v2) {
+        return new LongTuple<>(v1, v2);
+    }
+
+    private final T v1;
+    private final long v2;
+
+    private LongTuple(final T v1, final long v2) {
+        this.v1 = v1;
+        this.v2 = v2;
+    }
+
+    public T v1() {
+        return v1;
+    }
+
+    public long v2() {
+        return v2;
+    }
+
+    @Override
+    public boolean equals(final Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        LongTuple tuple = (LongTuple) o;
+
+        return (v1 == null ? tuple.v1 == null : v1.equals(tuple.v1)) && (v2 == tuple.v2);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = v1 != null ? v1.hashCode() : 0;
+        result = 31 * result + Long.hashCode(v2);
+        return result;
+    }
+
+    @Override
+    public String toString() {
+        return "Tuple [v1=" + v1 + ", v2=" + v2 + "]";
+    }
+
+}

core/src/main/java/org/elasticsearch/index/seqno/GlobalCheckpointTracker.java

@@ -21,6 +21,7 @@
 
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.shard.AbstractIndexShardComponent;
+import org.elasticsearch.index.shard.PrimaryContext;
 import org.elasticsearch.index.shard.ShardId;
 
 import java.util.Set;
@@ -165,13 +166,24 @@ public void updateGlobalCheckpointOnReplica(final long globalCheckpoint) {
 
     /**
      * Notifies the service of the current allocation IDs in the cluster state. See
-     * {@link GlobalCheckpointTracker#updateAllocationIdsFromMaster(Set, Set)} for details.
+     * {@link GlobalCheckpointTracker#updateAllocationIdsFromMaster(long, Set, Set)} for details.
      *
-     * @param activeAllocationIds       the allocation IDs of the currently active shard copies
-     * @param initializingAllocationIds the allocation IDs of the currently initializing shard copies
+     * @param applyingClusterStateVersion the cluster state version being applied when updating the allocation IDs from the master
+     * @param activeAllocationIds         the allocation IDs of the currently active shard copies
+     * @param initializingAllocationIds   the allocation IDs of the currently initializing shard copies
      */
-    public void updateAllocationIdsFromMaster(final Set<String> activeAllocationIds, final Set<String> initializingAllocationIds) {
-        globalCheckpointTracker.updateAllocationIdsFromMaster(activeAllocationIds, initializingAllocationIds);
+    public void updateAllocationIdsFromMaster(
+            final long applyingClusterStateVersion, final Set<String> activeAllocationIds, final Set<String> initializingAllocationIds) {
+        globalCheckpointTracker.updateAllocationIdsFromMaster(applyingClusterStateVersion, activeAllocationIds, initializingAllocationIds);
+    }
+
+    /**
+     * Updates the known allocation IDs and the local checkpoints for the corresponding allocations from a primary relocation source.
+     *
+     * @param primaryContext the sequence number context
+     */
+    public void updateAllocationIdsFromPrimaryContext(final PrimaryContext primaryContext) {
+        globalCheckpointTracker.updateAllocationIdsFromPrimaryContext(primaryContext);
     }
 
     /**
@@ -183,4 +195,20 @@ public boolean pendingInSync() {
         return globalCheckpointTracker.pendingInSync();
     }
 
+    /**
+     * Get the primary context for the shard. This includes the state of the global checkpoint tracker.
+     *
+     * @return the primary context
+     */
+    public PrimaryContext primaryContext() {
+        return globalCheckpointTracker.primaryContext();
+    }
+
+    /**
+     * Releases a previously acquired primary context.
+     */
+    public void releasePrimaryContext() {
+        globalCheckpointTracker.releasePrimaryContext();
+    }
+
 }

core/src/main/java/org/elasticsearch/index/seqno/SequenceNumbersService.java

@@ -515,31 +515,37 @@ public IndexShardState markAsRecovering(String reason, RecoveryState recoverySta
 
     private final AtomicBoolean primaryReplicaResyncInProgress = new AtomicBoolean();
 
-    public void relocated(String reason) throws IllegalIndexShardStateException, InterruptedException {
+    /**
+     * Completes the relocation. Operations are blocked and current operations are drained before changing state to relocated. The provided
+     * {@link Runnable} is executed after all operations are successfully blocked.
+     *
+     * @param reason    the reason for the relocation
+     * @param consumer a {@link Runnable} that is executed after operations are blocked
+     * @throws IllegalIndexShardStateException if the shard is not relocating due to concurrent cancellation
+     * @throws InterruptedException            if blocking operations is interrupted
+     */
+    public void relocated(
+            final String reason, final Consumer<PrimaryContext> consumer) throws IllegalIndexShardStateException, InterruptedException {
         assert shardRouting.primary() : "only primaries can be marked as relocated: " + shardRouting;
         try {
             indexShardOperationPermits.blockOperations(30, TimeUnit.MINUTES, () -> {
                 // no shard operation permits are being held here, move state from started to relocated
                 assert indexShardOperationPermits.getActiveOperationsCount() == 0 :
-                    "in-flight operations in progress while moving shard state to relocated";
-                synchronized (mutex) {
-                    if (state != IndexShardState.STARTED) {
-                        throw new IndexShardNotStartedException(shardId, state);
-                    }
-                    // if the master cancelled the recovery, the target will be removed
-                    // and the recovery will stopped.
-                    // However, it is still possible that we concurrently end up here
-                    // and therefore have to protect we don't mark the shard as relocated when
-                    // its shard routing says otherwise.
-                    if (shardRouting.relocating() == false) {
-                        throw new IllegalIndexShardStateException(shardId, IndexShardState.STARTED,
-                            ": shard is no longer relocating " + shardRouting);
-                    }
-                    if (primaryReplicaResyncInProgress.get()) {
-                        throw new IllegalIndexShardStateException(shardId, IndexShardState.STARTED,
-                            ": primary relocation is forbidden while primary-replica resync is in progress " + shardRouting);
+                        "in-flight operations in progress while moving shard state to relocated";
+                /*
+                 * We should not invoke the runnable under the mutex as the expected implementation is to handoff the primary context via a
+                 * network operation. Doing this under the mutex can implicitly block the cluster state update thread on network operations.
+                 */
+                verifyRelocatingState();
+                final PrimaryContext primaryContext = getEngine().seqNoService().primaryContext();
+                try {
+                    consumer.accept(primaryContext);
+                    synchronized (mutex) {
+                        verifyRelocatingState();
+                        changeState(IndexShardState.RELOCATED, reason);
                     }
-                    changeState(IndexShardState.RELOCATED, reason);
+                } catch (final Exception e) {
+                    getEngine().seqNoService().releasePrimaryContext();
                 }
             });
         } catch (TimeoutException e) {
@@ -551,6 +557,26 @@ public void relocated(String reason) throws IllegalIndexShardStateException, Int
         }
     }
 
+    private void verifyRelocatingState() {
+        if (state != IndexShardState.STARTED) {
+            throw new IndexShardNotStartedException(shardId, state);
+        }
+        /*
+         * If the master cancelled recovery, the target will be removed and the recovery will be cancelled. However, it is still possible
+         * that we concurrently end up here and therefore have to protect that we do not mark the shard as relocated when its shard routing
+         * says otherwise.
+         */
+
+        if (shardRouting.relocating() == false) {
+            throw new IllegalIndexShardStateException(shardId, IndexShardState.STARTED,
+                ": shard is no longer relocating " + shardRouting);
+        }
+
+        if (primaryReplicaResyncInProgress.get()) {
+            throw new IllegalIndexShardStateException(shardId, IndexShardState.STARTED,
+                ": primary relocation is forbidden while primary-replica resync is in progress " + shardRouting);
+        }
+    }
 
     public IndexShardState state() {
         return state;
@@ -1319,16 +1345,16 @@ private void ensureWriteAllowed(Engine.Operation.Origin origin) throws IllegalIn
 
     private void verifyPrimary() {
         if (shardRouting.primary() == false) {
-            throw new IllegalStateException("shard is not a primary " + shardRouting);
+            throw new IllegalStateException("shard " + shardRouting + " is not a primary");
         }
     }
 
     private void verifyReplicationTarget() {
         final IndexShardState state = state();
         if (shardRouting.primary() && shardRouting.active() && state != IndexShardState.RELOCATED) {
             // must use exception that is not ignored by replication logic. See TransportActions.isShardNotAvailableException
-            throw new IllegalStateException("active primary shard cannot be a replication target before " +
-                " relocation hand off " + shardRouting + ", state is [" + state + "]");
+            throw new IllegalStateException("active primary shard " + shardRouting + " cannot be a replication target before " +
+                "relocation hand off, state is [" + state + "]");
         }
     }
 
@@ -1603,8 +1629,8 @@ public void markAllocationIdAsInSync(final String allocationId, final long local
         verifyPrimary();
         getEngine().seqNoService().markAllocationIdAsInSync(allocationId, localCheckpoint);
         /*
-         * We could have blocked waiting for the replica to catch up that we fell idle and there will not be a background sync to the
-         * replica; mark our self as active to force a future background sync.
+         * We could have blocked so long waiting for the replica to catch up that we fell idle and there will not be a background sync to
+         * the replica; mark our self as active to force a future background sync.
          */
         active.compareAndSet(false, true);
     }
@@ -1654,18 +1680,34 @@ public void updateGlobalCheckpointOnReplica(final long globalCheckpoint) {
 
     /**
      * Notifies the service of the current allocation IDs in the cluster state. See
-     * {@link org.elasticsearch.index.seqno.GlobalCheckpointTracker#updateAllocationIdsFromMaster(Set, Set)}
+     * {@link org.elasticsearch.index.seqno.GlobalCheckpointTracker#updateAllocationIdsFromMaster(long, Set, Set)}
      * for details.
      *
-     * @param activeAllocationIds       the allocation IDs of the currently active shard copies
-     * @param initializingAllocationIds the allocation IDs of the currently initializing shard copies
+     * @param applyingClusterStateVersion the cluster state version being applied when updating the allocation IDs from the master
+     * @param activeAllocationIds         the allocation IDs of the currently active shard copies
+     * @param initializingAllocationIds   the allocation IDs of the currently initializing shard copies
      */
-    public void updateAllocationIdsFromMaster(final Set<String> activeAllocationIds, final Set<String> initializingAllocationIds) {
+    public void updateAllocationIdsFromMaster(
+            final long applyingClusterStateVersion, final Set<String> activeAllocationIds, final Set<String> initializingAllocationIds) {
         verifyPrimary();
         final Engine engine = getEngineOrNull();
         // if the engine is not yet started, we are not ready yet and can just ignore this
         if (engine != null) {
-            engine.seqNoService().updateAllocationIdsFromMaster(activeAllocationIds, initializingAllocationIds);
+            engine.seqNoService().updateAllocationIdsFromMaster(applyingClusterStateVersion, activeAllocationIds, initializingAllocationIds);
+        }
+    }
+
+    /**
+     * Updates the known allocation IDs and the local checkpoints for the corresponding allocations from a primary relocation source.
+     *
+     * @param primaryContext the sequence number context
+     */
+    public void updateAllocationIdsFromPrimaryContext(final PrimaryContext primaryContext) {
+        verifyPrimary();
+        assert shardRouting.isRelocationTarget() : "only relocation target can update allocation IDs from primary context: " + shardRouting;
+        final Engine engine = getEngineOrNull();
+        if (engine != null) {
+            engine.seqNoService().updateAllocationIdsFromPrimaryContext(primaryContext);
         }
     }
 

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

@@ -0,0 +1,105 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.index.shard;
+
+import com.carrotsearch.hppc.ObjectLongHashMap;
+import com.carrotsearch.hppc.ObjectLongMap;
+import com.carrotsearch.hppc.cursors.ObjectLongCursor;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+
+import java.io.IOException;
+
+/**
+ * Represents the sequence number component of the primary context. This is the knowledge on the primary of the in-sync and initializing
+ * shards and their local checkpoints.
+ */
+public class PrimaryContext implements Writeable {
+
+    private long clusterStateVersion;
+
+    public long clusterStateVersion() {
+        return clusterStateVersion;
+    }
+
+    private ObjectLongMap<String> inSyncLocalCheckpoints;
+
+    public ObjectLongMap<String> inSyncLocalCheckpoints() {
+        return inSyncLocalCheckpoints;
+    }
+
+    private ObjectLongMap<String> trackingLocalCheckpoints;
+
+    public ObjectLongMap<String> trackingLocalCheckpoints() {
+        return trackingLocalCheckpoints;
+    }
+
+    public PrimaryContext(
+            final long clusterStateVersion,
+            final ObjectLongMap<String> inSyncLocalCheckpoints,
+            final ObjectLongMap<String> trackingLocalCheckpoints) {
+        this.clusterStateVersion = clusterStateVersion;
+        this.inSyncLocalCheckpoints = inSyncLocalCheckpoints;
+        this.trackingLocalCheckpoints = trackingLocalCheckpoints;
+    }
+
+    public PrimaryContext(final StreamInput in) throws IOException {
+        clusterStateVersion = in.readVLong();
+        inSyncLocalCheckpoints = readMap(in);
+        trackingLocalCheckpoints = readMap(in);
+    }
+
+    private static ObjectLongMap<String> readMap(final StreamInput in) throws IOException {
+        final int length = in.readVInt();
+        final ObjectLongMap<String> map = new ObjectLongHashMap<>(length);
+        for (int i = 0; i < length; i++) {
+            final String key = in.readString();
+            final long value = in.readZLong();
+            map.addTo(key, value);
+        }
+        return map;
+    }
+
+    @Override
+    public void writeTo(final StreamOutput out) throws IOException {
+        out.writeVLong(clusterStateVersion);
+        writeMap(out, inSyncLocalCheckpoints);
+        writeMap(out, trackingLocalCheckpoints);
+    }
+
+    private static void writeMap(final StreamOutput out, final ObjectLongMap<String> map) throws IOException {
+        out.writeVInt(map.size());
+        for (ObjectLongCursor<String> cursor : map) {
+            out.writeString(cursor.key);
+            out.writeZLong(cursor.value);
+        }
+    }
+
+    @Override
+    public String toString() {
+        return "PrimaryContext{" +
+                "clusterStateVersion=" + clusterStateVersion +
+                ", inSyncLocalCheckpoints=" + inSyncLocalCheckpoints +
+                ", trackingLocalCheckpoints=" + trackingLocalCheckpoints +
+                '}';
+    }
+
+}