Allow tracking evacuated IDs since the last refresh in LiveVersionMap (#95331)

This PR adds adds an interface and noop implementation for LiveVersionMapArchive.
It receives old maps from the LiveVersionMap upon a refresh, and is also informed
when unpromotable shards have been refreshed (one way of doing this is implemented
in the matching Serverless PR).

relates ES-5728

Author: pxsalehi

server/src/main/java/org/elasticsearch/index/engine/InternalEngine.java

@@ -141,7 +141,8 @@ public class InternalEngine extends Engine {
 
     // A uid (in the form of BytesRef) to the version map
     // we use the hashed variant since we iterate over it and check removal and additions on existing keys
-    private final LiveVersionMap versionMap = new LiveVersionMap();
+    private final LiveVersionMap versionMap;
+    private final LiveVersionMapArchive liveVersionMapArchive;
 
     private volatile SegmentInfos lastCommittedSegmentInfos;
 
@@ -216,6 +217,8 @@ public InternalEngine(EngineConfig engineConfig) {
         this.maxDocs = maxDocs;
         this.relativeTimeInNanosSupplier = config().getRelativeTimeInNanosSupplier();
         this.lastFlushTimestamp = relativeTimeInNanosSupplier.getAsLong(); // default to creation timestamp
+        this.liveVersionMapArchive = createLiveVersionMapArchive();
+        this.versionMap = new LiveVersionMap(liveVersionMapArchive);
         final TranslogDeletionPolicy translogDeletionPolicy = new TranslogDeletionPolicy();
         store.incRef();
         IndexWriter writer = null;
@@ -3170,4 +3173,17 @@ protected void waitForCommitDurability(long generation, ActionListener<Void> lis
             listener.onResponse(null);
         }
     }
+
+    protected LiveVersionMapArchive createLiveVersionMapArchive() {
+        return LiveVersionMapArchive.NOOP_ARCHIVE;
+    }
+
+    protected LiveVersionMapArchive getLiveVersionMapArchive() {
+        return liveVersionMapArchive;
+    }
+
+    // Visible for testing purposes only
+    public LiveVersionMap getLiveVersionMap() {
+        return versionMap;
+    }
 }

server/src/main/java/org/elasticsearch/index/engine/LiveVersionMap.java

@@ -22,11 +22,21 @@
 import java.util.concurrent.atomic.AtomicLong;
 
 /** Maps _uid value to its version information. */
-final class LiveVersionMap implements ReferenceManager.RefreshListener, Accountable {
+public final class LiveVersionMap implements ReferenceManager.RefreshListener, Accountable {
 
     private final KeyedLock<BytesRef> keyedLock = new KeyedLock<>();
 
-    private static final class VersionLookup {
+    private final LiveVersionMapArchive archiver;
+
+    LiveVersionMap() {
+        this(LiveVersionMapArchive.NOOP_ARCHIVE);
+    }
+
+    LiveVersionMap(LiveVersionMapArchive archiver) {
+        this.archiver = archiver;
+    }
+
+    public static final class VersionLookup {
 
         /** Tracks bytes used by current map, i.e. what is freed on refresh. For deletes, which are also added to tombstones,
          *  we only account for the CHM entry here, and account for BytesRef/VersionValue against the tombstones, since refresh would not
@@ -52,19 +62,25 @@ private static final class VersionLookup {
         // the tombstone
         private final AtomicLong minDeleteTimestamp = new AtomicLong(Long.MAX_VALUE);
 
+        // Modifies the map of this instance by merging with the given VersionLookup
+        public void merge(VersionLookup versionLookup) {
+            map.putAll(versionLookup.map);
+            minDeleteTimestamp.accumulateAndGet(versionLookup.minDeleteTimestamp(), Math::min);
+        }
+
         private VersionLookup(Map<BytesRef, VersionValue> map) {
             this.map = map;
         }
 
-        VersionValue get(BytesRef key) {
+        public VersionValue get(BytesRef key) {
             return map.get(key);
         }
 
         VersionValue put(BytesRef key, VersionValue value) {
             return map.put(key, value);
         }
 
-        boolean isEmpty() {
+        public boolean isEmpty() {
             return map.isEmpty();
         }
 
@@ -88,6 +104,9 @@ public void updateMinDeletedTimestamp(DeleteVersionValue delete) {
             minDeleteTimestamp.accumulateAndGet(delete.time, Math::min);
         }
 
+        public long minDeleteTimestamp() {
+            return minDeleteTimestamp.get();
+        }
     }
 
     private static final class Maps {
@@ -98,7 +117,7 @@ private static final class Maps {
         // Used while refresh is running, and to hold adds/deletes until refresh finishes. We read from both current and old on lookup:
         final VersionLookup old;
 
-        // this is not volatile since we don't need to maintain a happens before relation ship across doc IDs so it's enough to
+        // this is not volatile since we don't need to maintain a happens before relationship across doc IDs so it's enough to
         // have the volatile read of the Maps reference to make it visible even across threads.
         boolean needsSafeAccess;
         final boolean previousMapsNeededSafeAccess;
@@ -135,10 +154,18 @@ Maps buildTransitionMap() {
             );
         }
 
+        /**
+         * similar to `invalidateOldMap` but used only for the `unsafeKeysMap` used for assertions
+         */
+        Maps invalidateOldMapForAssert() {
+            return new Maps(current, VersionLookup.EMPTY, previousMapsNeededSafeAccess);
+        }
+
         /**
          * builds a new map that invalidates the old map but maintains the current. This should be called in afterRefresh()
          */
-        Maps invalidateOldMap() {
+        Maps invalidateOldMap(LiveVersionMapArchive archiver) {
+            archiver.afterRefresh(old);
             return new Maps(current, VersionLookup.EMPTY, previousMapsNeededSafeAccess);
         }
 
@@ -245,8 +272,8 @@ public void afterRefresh(boolean didRefresh) throws IOException {
         // reopen, and so any concurrent indexing requests can still sneak in a few additions to that current map that are in fact
         // reflected in the previous reader. We don't touch tombstones here: they expire on their own index.gc_deletes timeframe:
 
-        maps = maps.invalidateOldMap();
-        assert (unsafeKeysMap = unsafeKeysMap.invalidateOldMap()) != null;
+        maps = maps.invalidateOldMap(archiver);
+        assert (unsafeKeysMap = unsafeKeysMap.invalidateOldMapForAssert()) != null;
 
     }
 
@@ -270,7 +297,14 @@ private VersionValue getUnderLock(final BytesRef uid, Maps currentMaps) {
             return value;
         }
 
-        return tombstones.get(uid);
+        // We first check the tombstone then the archiver since the archiver accumulates ids from the old map, and we
+        // makes sure in `putDeleteUnderLock` the old map does not hold an entry that is in tombstone, archiver also wouldn't have them.
+        value = tombstones.get(uid);
+        if (value != null) {
+            return value;
+        }
+
+        return archiver.get(uid);
     }
 
     VersionValue getVersionForAssert(final BytesRef uid) {
@@ -368,7 +402,8 @@ private boolean canRemoveTombstone(long maxTimestampToPrune, long maxSeqNoToPrun
         // version value can't be removed it's
         // not yet flushed to lucene ie. it's part of this current maps object
         final boolean isNotTrackedByCurrentMaps = versionValue.time < maps.getMinDeleteTimestamp();
-        return isTooOld && isSafeToPrune && isNotTrackedByCurrentMaps;
+        final boolean isNotTrackedByArchiver = versionValue.time < archiver.getMinDeleteTimestamp();
+        return isTooOld && isSafeToPrune && isNotTrackedByCurrentMaps & isNotTrackedByArchiver;
     }
 
     /**
@@ -459,4 +494,9 @@ boolean assertKeyedLockHeldByCurrentThread(BytesRef uid) {
         assert keyedLock.isHeldByCurrentThread(uid) : "Thread [" + Thread.currentThread().getName() + "], uid [" + uid.utf8ToString() + "]";
         return true;
     }
+
+    // visible for testing purposes only
+    LiveVersionMapArchive getArchiver() {
+        return archiver;
+    }
 }

server/src/main/java/org/elasticsearch/index/engine/LiveVersionMapArchive.java

@@ -0,0 +1,48 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.index.engine;
+
+import org.apache.lucene.util.BytesRef;
+
+/**
+ * Keeps track of the old map of a LiveVersionMap that gets evacuated on a refresh
+ */
+public interface LiveVersionMapArchive {
+    /**
+     * Archive the old map evacuated due to a refresh
+     *
+     * @param old is the old map that is evacuated on a refresh
+     */
+    void afterRefresh(LiveVersionMap.VersionLookup old);
+
+    /**
+     * Look up the given uid in the archive
+     */
+    VersionValue get(BytesRef uid);
+
+    /**
+     * Returns the min delete timestamp across all archived maps.
+     */
+    long getMinDeleteTimestamp();
+
+    LiveVersionMapArchive NOOP_ARCHIVE = new LiveVersionMapArchive() {
+        @Override
+        public void afterRefresh(LiveVersionMap.VersionLookup old) {}
+
+        @Override
+        public VersionValue get(BytesRef uid) {
+            return null;
+        }
+
+        @Override
+        public long getMinDeleteTimestamp() {
+            return Long.MAX_VALUE;
+        }
+    };
+}

server/src/main/java/org/elasticsearch/index/engine/VersionValue.java

@@ -13,7 +13,7 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.translog.Translog;
 
-abstract class VersionValue implements Accountable {
+public abstract class VersionValue implements Accountable {
 
     private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(VersionValue.class);
 

server/src/test/java/org/elasticsearch/index/engine/LiveVersionMapArchiveTests.java

@@ -0,0 +1,76 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.index.engine;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.index.shard.IndexShard;
+import org.elasticsearch.index.shard.IndexShardTestCase;
+import org.hamcrest.CoreMatchers;
+import org.hamcrest.core.IsInstanceOf;
+
+import static org.elasticsearch.index.engine.LiveVersionMapTestUtils.putIndex;
+import static org.elasticsearch.index.engine.LiveVersionMapTestUtils.randomIndexVersionValue;
+
+public class LiveVersionMapArchiveTests extends IndexShardTestCase {
+    private class TestArchive implements LiveVersionMapArchive {
+
+        boolean afterRefreshCalled = false;
+        LiveVersionMap.VersionLookup archivedMap = null;
+
+        @Override
+        public void afterRefresh(LiveVersionMap.VersionLookup old) {
+            afterRefreshCalled = true;
+            archivedMap = old;
+        }
+
+        @Override
+        public VersionValue get(BytesRef uid) {
+            return null;
+        }
+
+        @Override
+        public long getMinDeleteTimestamp() {
+            return Long.MAX_VALUE;
+        }
+    }
+
+    private class TestEngine extends InternalEngine {
+        TestEngine(EngineConfig engineConfig) {
+            super(engineConfig);
+        }
+
+        @Override
+        public LiveVersionMapArchive createLiveVersionMapArchive() {
+            return new TestArchive();
+        }
+    }
+
+    public void testLiveVersionMapArchiveCreation() throws Exception {
+        final IndexShard shard = newStartedShard(false, Settings.EMPTY, TestEngine::new);
+        assertThat(shard.getEngineOrNull(), CoreMatchers.instanceOf(InternalEngine.class));
+        var engine = (InternalEngine) shard.getEngineOrNull();
+        assertThat(engine.getLiveVersionMapArchive(), IsInstanceOf.instanceOf(TestArchive.class));
+        closeShards(shard);
+    }
+
+    public void testLiveVersionMapArchiveBasic() throws Exception {
+        TestArchive archive = new TestArchive();
+        LiveVersionMap map = new LiveVersionMap(archive);
+        putIndex(map, "1", randomIndexVersionValue());
+        map.beforeRefresh();
+        assertFalse(archive.afterRefreshCalled);
+        assertNull(archive.archivedMap);
+        map.afterRefresh(randomBoolean());
+        assertTrue(archive.afterRefreshCalled);
+        var archived = archive.archivedMap;
+        assertNotNull(archived);
+        assertEquals(archived.size(), 1);
+    }
+}

server/src/test/java/org/elasticsearch/index/engine/LiveVersionMapTestUtils.java

@@ -0,0 +1,81 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.index.engine;
+
+import org.apache.lucene.index.Term;
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.index.mapper.IdFieldMapper;
+import org.elasticsearch.index.mapper.Uid;
+import org.elasticsearch.index.translog.Translog;
+
+import static org.elasticsearch.test.ESTestCase.randomBoolean;
+import static org.elasticsearch.test.ESTestCase.randomInt;
+import static org.elasticsearch.test.ESTestCase.randomNonNegativeLong;
+
+public class LiveVersionMapTestUtils {
+
+    public static LiveVersionMap newLiveVersionMap(LiveVersionMapArchive archiver) {
+        return new LiveVersionMap(archiver);
+    }
+
+    public static DeleteVersionValue newDeleteVersionValue(long version, long seqNo, long term, long time) {
+        return new DeleteVersionValue(version, seqNo, term, time);
+    }
+
+    public static IndexVersionValue newIndexVersionValue(Translog.Location location, long version, long seqNo, long term) {
+        return new IndexVersionValue(location, version, seqNo, term);
+    }
+
+    public static VersionValue get(LiveVersionMap map, String id) {
+        try (Releasable r = acquireLock(map, uid(id))) {
+            return map.getUnderLock(uid(id));
+        }
+    }
+
+    public static void putIndex(LiveVersionMap map, String id, IndexVersionValue version) {
+        try (Releasable r = acquireLock(map, uid(id))) {
+            map.putIndexUnderLock(uid(id), version);
+        }
+    }
+
+    public static void putDelete(LiveVersionMap map, String id, DeleteVersionValue version) {
+        try (Releasable r = acquireLock(map, uid(id))) {
+            map.putDeleteUnderLock(uid(id), version);
+        }
+    }
+
+    public static void pruneTombstones(LiveVersionMap map, long maxTimestampToPrune, long maxSeqNoToPrune) {
+        map.pruneTombstones(maxTimestampToPrune, maxSeqNoToPrune);
+    }
+
+    static IndexVersionValue randomIndexVersionValue() {
+        return new IndexVersionValue(randomTranslogLocation(), randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong());
+    }
+
+    static Translog.Location randomTranslogLocation() {
+        if (randomBoolean()) {
+            return null;
+        } else {
+            return new Translog.Location(randomNonNegativeLong(), randomNonNegativeLong(), randomInt());
+        }
+    }
+
+    public static int VersionLookupSize(LiveVersionMap.VersionLookup lookup) {
+        return lookup.size();
+    }
+
+    private static Releasable acquireLock(LiveVersionMap map, BytesRef uid) {
+        return map.acquireLock(uid);
+    }
+
+    public static BytesRef uid(String id) {
+        return new Term(IdFieldMapper.NAME, Uid.encodeId(id)).bytes();
+    }
+}