Update shardGenerations for all indices on snapshot finalization

docs/changelog/128650.yaml

@@ -0,0 +1,5 @@
+pr: 128650
+summary: Always update `shardGenerations` for previous in-progress snapshots
+area: Snapshot/Restore
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/repositories/FinalizeSnapshotContext.java

@@ -27,7 +27,7 @@
  */
 public final class FinalizeSnapshotContext extends DelegatingActionListener<RepositoryData, RepositoryData> {
 
-    private final ShardGenerations updatedShardGenerations;
+    private final UpdatedShardGenerations updatedShardGenerations;
 
     /**
      * Obsolete shard generations map computed from the cluster state update that this finalization executed in
@@ -46,7 +46,7 @@ public final class FinalizeSnapshotContext extends DelegatingActionListener<Repo
     private final Runnable onDone;
 
     /**
-     * @param updatedShardGenerations updated shard generations
+     * @param updatedShardGenerations updated shard generations for both live and deleted indices
      * @param repositoryStateId       the unique id identifying the state of the repository when the snapshot began
      * @param clusterMetadata         cluster metadata
      * @param snapshotInfo            SnapshotInfo instance to write for this snapshot
@@ -57,7 +57,7 @@ public final class FinalizeSnapshotContext extends DelegatingActionListener<Repo
      *                                once all cleanup operations after snapshot completion have executed
      */
     public FinalizeSnapshotContext(
-        ShardGenerations updatedShardGenerations,
+        UpdatedShardGenerations updatedShardGenerations,
         long repositoryStateId,
         Metadata clusterMetadata,
         SnapshotInfo snapshotInfo,
@@ -78,7 +78,7 @@ public long repositoryStateId() {
         return repositoryStateId;
     }
 
-    public ShardGenerations updatedShardGenerations() {
+    public UpdatedShardGenerations updatedShardGenerations() {
         return updatedShardGenerations;
     }
 
@@ -120,4 +120,24 @@ public void onDone() {
     public void onResponse(RepositoryData repositoryData) {
         delegate.onResponse(repositoryData);
     }
+
+    /**
+     * A record used to track the new shard generations that have been written for each shard in a snapshot.
+     * An index may be deleted after the shard generation is written but before the snapshot is finalized.
+     * In this case, its shard generation is tracked in {@link #deletedIndices} because it's still a valid
+     * shard generation blob that exists in the repository and may be used by subsequent snapshots, even though
+     * the index will not be included in the snapshot being finalized. Otherwise, it is tracked in
+     * {@link #liveIndices}.
+     */
+    public record UpdatedShardGenerations(ShardGenerations liveIndices, ShardGenerations deletedIndices) {
+        public static final UpdatedShardGenerations EMPTY = new UpdatedShardGenerations(ShardGenerations.EMPTY, ShardGenerations.EMPTY);
+
+        public UpdatedShardGenerations(ShardGenerations updated) {
+            this(updated, ShardGenerations.EMPTY);
+        }
+
+        public boolean hasShardGen(RepositoryShardId repositoryShardId) {
+            return liveIndices.hasShardGen(repositoryShardId) || deletedIndices.hasShardGen(repositoryShardId);
+        }
+    }
 }

server/src/main/java/org/elasticsearch/repositories/RepositoryData.java

@@ -25,6 +25,7 @@
 import org.elasticsearch.index.IndexVersions;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
+import org.elasticsearch.repositories.FinalizeSnapshotContext.UpdatedShardGenerations;
 import org.elasticsearch.snapshots.SnapshotId;
 import org.elasticsearch.snapshots.SnapshotInfo;
 import org.elasticsearch.snapshots.SnapshotState;
@@ -405,16 +406,16 @@ public Map<IndexId, Collection<String>> indexMetaDataToRemoveAfterRemovingSnapsh
      *
      * @param snapshotId       Id of the new snapshot
      * @param details          Details of the new snapshot
-     * @param shardGenerations Updated shard generations in the new snapshot. For each index contained in the snapshot an array of new
-     *                         generations indexed by the shard id they correspond to must be supplied.
+     * @param updatedShardGenerations Updated shard generations in the new snapshot, including both indices that are included
+     *                                in the given snapshot and those got deleted while finalizing.
      * @param indexMetaBlobs   Map of index metadata blob uuids
      * @param newIdentifiers   Map of new index metadata blob uuids keyed by the identifiers of the
      *                         {@link IndexMetadata} in them
      */
     public RepositoryData addSnapshot(
         final SnapshotId snapshotId,
         final SnapshotDetails details,
-        final ShardGenerations shardGenerations,
+        final UpdatedShardGenerations updatedShardGenerations,
         @Nullable final Map<IndexId, String> indexMetaBlobs,
         @Nullable final Map<String, String> newIdentifiers
     ) {
@@ -424,6 +425,7 @@ public RepositoryData addSnapshot(
             // the new master, so we make the operation idempotent
             return this;
         }
+        final var shardGenerations = updatedShardGenerations.liveIndices();
         Map<String, SnapshotId> snapshots = new HashMap<>(snapshotIds);
         snapshots.put(snapshotId.getUUID(), snapshotId);
         Map<String, SnapshotDetails> newSnapshotDetails = new HashMap<>(snapshotsDetails);
@@ -459,7 +461,7 @@ public RepositoryData addSnapshot(
             snapshots,
             newSnapshotDetails,
             allIndexSnapshots,
-            ShardGenerations.builder().putAll(this.shardGenerations).putAll(shardGenerations).build(),
+            ShardGenerations.builder().putAll(this.shardGenerations).update(updatedShardGenerations).build(),
             newIndexMetaGenerations,
             clusterUUID
         );

server/src/main/java/org/elasticsearch/repositories/ShardGenerations.java

@@ -28,6 +28,8 @@
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
+import static org.elasticsearch.repositories.FinalizeSnapshotContext.UpdatedShardGenerations;
+
 /**
  * Represents the current {@link ShardGeneration} for each shard in a repository.
  */
@@ -231,6 +233,14 @@ public Builder putAll(ShardGenerations shardGenerations) {
             return this;
         }
 
+        public Builder update(UpdatedShardGenerations updatedShardGenerations) {
+            putAll(updatedShardGenerations.liveIndices());
+            // For deleted indices, we only update the generations if they are present in the existing generations, i.e.
+            // they are referenced by other snapshots.
+            updateIfPresent(updatedShardGenerations.deletedIndices());
+            return this;
+        }
+
         public Builder put(IndexId indexId, int shardId, SnapshotsInProgress.ShardSnapshotStatus status) {
             // only track generations for successful shard status values
             return put(indexId, shardId, status.state().failed() ? null : status.generation());
@@ -244,6 +254,20 @@ public Builder put(IndexId indexId, int shardId, ShardGeneration generation) {
             return this;
         }
 
+        private Builder updateIfPresent(ShardGenerations shardGenerations) {
+            shardGenerations.shardGenerations.forEach((indexId, gens) -> {
+                if (generations.containsKey(indexId)) {
+                    for (int i = 0; i < gens.size(); i++) {
+                        final ShardGeneration gen = gens.get(i);
+                        if (gen != null) {
+                            generations.get(indexId).put(i, gen);
+                        }
+                    }
+                }
+            });
+            return this;
+        }
+
         private boolean noDuplicateIndicesWithSameName(IndexId newId) {
             for (IndexId id : generations.keySet()) {
                 if (id.getName().equals(newId.getName()) && id.equals(newId) == false) {

server/src/main/java/org/elasticsearch/repositories/blobstore/BlobStoreRepository.java

@@ -1749,7 +1749,7 @@ int sizeInBytes() {
     public void finalizeSnapshot(final FinalizeSnapshotContext finalizeSnapshotContext) {
         assert ThreadPool.assertCurrentThreadPool(ThreadPool.Names.SNAPSHOT);
         final long repositoryStateId = finalizeSnapshotContext.repositoryStateId();
-        final ShardGenerations shardGenerations = finalizeSnapshotContext.updatedShardGenerations();
+        final ShardGenerations shardGenerations = finalizeSnapshotContext.updatedShardGenerations().liveIndices();
         final SnapshotInfo snapshotInfo = finalizeSnapshotContext.snapshotInfo();
         assert repositoryStateId > RepositoryData.UNKNOWN_REPO_GEN
             : "Must finalize based on a valid repository generation but received [" + repositoryStateId + "]";
@@ -1867,7 +1867,7 @@ record RootBlobUpdateResult(RepositoryData oldRepositoryData, RepositoryData new
                     existingRepositoryData.addSnapshot(
                         snapshotId,
                         snapshotDetails,
-                        shardGenerations,
+                        finalizeSnapshotContext.updatedShardGenerations(),
                         metadataWriteResult.indexMetas(),
                         metadataWriteResult.indexMetaIdentifiers()
                     ),