From 4bc49112354258e7650abcb4897442ebaf91473a Mon Sep 17 00:00:00 2001
From: David Turner <david.turner@elastic.co>
Date: Mon, 27 Nov 2023 03:01:12 -0800
Subject: [PATCH] Assert aborted snapshots have no pending shard work (#102621)

When a snapshot is aborted, all shard snapshots must either be complete
(`SUCCESS`, `FAILED` or `MISSING`) or else they must be in state
`ABORTED` while the data node finishes its work. It seems we do not
check this invariant anywhere today, so this commit adds an assertion of
this invariant.
---
 .../elasticsearch/cluster/SnapshotsInProgress.java  |  8 +++++++-
 .../SnapshotsInProgressSerializationTests.java      | 13 ++++++++++---
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/cluster/SnapshotsInProgress.java b/server/src/main/java/org/elasticsearch/cluster/SnapshotsInProgress.java
index 34a0024c1a467..1a079d03405d7 100644
--- a/server/src/main/java/org/elasticsearch/cluster/SnapshotsInProgress.java
+++ b/server/src/main/java/org/elasticsearch/cluster/SnapshotsInProgress.java
@@ -343,14 +343,20 @@ private static boolean assertConsistentEntries(Map<String, ByRepo> entries) {
                 assert entry.repository().equals(repository) : "mismatched repository " + entry + " tracked under " + repository;
                 for (Map.Entry<RepositoryShardId, ShardSnapshotStatus> shard : entry.shardsByRepoShardId().entrySet()) {
                     final RepositoryShardId sid = shard.getKey();
+                    final ShardSnapshotStatus shardSnapshotStatus = shard.getValue();
                     assert assertShardStateConsistent(
                         entriesForRepository,
                         assignedShards,
                         queuedShards,
                         sid.indexName(),
                         sid.shardId(),
-                        shard.getValue()
+                        shardSnapshotStatus
                     );
+
+                    assert entry.state() != State.ABORTED
+                        || shardSnapshotStatus.state == ShardState.ABORTED
+                        || shardSnapshotStatus.state().completed()
+                        : sid + " is in state " + shardSnapshotStatus.state() + " in aborted snapshot " + entry.snapshot;
                 }
             }
             // make sure in-flight-shard-states can be built cleanly for the entries without tripping assertions
diff --git a/server/src/test/java/org/elasticsearch/snapshots/SnapshotsInProgressSerializationTests.java b/server/src/test/java/org/elasticsearch/snapshots/SnapshotsInProgressSerializationTests.java
index b5ab6be586670..74277954b8002 100644
--- a/server/src/test/java/org/elasticsearch/snapshots/SnapshotsInProgressSerializationTests.java
+++ b/server/src/test/java/org/elasticsearch/snapshots/SnapshotsInProgressSerializationTests.java
@@ -523,8 +523,15 @@ public void testXContent() throws IOException {
     }
 
     public static State randomState(Map<ShardId, SnapshotsInProgress.ShardSnapshotStatus> shards) {
-        return SnapshotsInProgress.completed(shards.values())
-            ? randomFrom(State.SUCCESS, State.FAILED)
-            : randomFrom(State.STARTED, State.INIT, State.ABORTED);
+        if (SnapshotsInProgress.completed(shards.values())) {
+            return randomFrom(State.SUCCESS, State.FAILED);
+        }
+        if (shards.values()
+            .stream()
+            .map(SnapshotsInProgress.ShardSnapshotStatus::state)
+            .allMatch(st -> st.completed() || st == ShardState.ABORTED)) {
+            return State.ABORTED;
+        }
+        return randomFrom(State.STARTED, State.INIT);
     }
 }