Introduce BlobStoreRepository CAS Mechanism (#93825)

Only for testing purposes through the `FsRepository` for now and rather simple,
but should get the job done and technically be correct for a compliant NFS implementation.

Co-authored-by: David Turner <david.turner@elastic.co>

modules/repository-azure/src/main/java/org/elasticsearch/repositories/azure/AzureBlobContainer.java

@@ -154,4 +154,19 @@ public class AzureBlobContainer extends AbstractBlobContainer {
     protected String buildKey(String blobName) {
         return keyPath + (blobName == null ? "" : blobName);
     }
+
+    @Override
+    public long compareAndExchangeRegister(String key, long expected, long updated) {
+        throw new UnsupportedOperationException(); // TODO
+    }
+
+    @Override
+    public boolean compareAndSetRegister(String key, long expected, long updated) {
+        throw new UnsupportedOperationException(); // TODO
+    }
+
+    @Override
+    public long getRegister(String key) throws IOException {
+        throw new UnsupportedOperationException(); // TODO
+    }
 }

modules/repository-gcs/src/main/java/org/elasticsearch/repositories/gcs/GoogleCloudStorageBlobContainer.java

@@ -117,4 +117,19 @@ class GoogleCloudStorageBlobContainer extends AbstractBlobContainer {
         assert blobName != null;
         return path + blobName;
     }
+
+    @Override
+    public long compareAndExchangeRegister(String key, long expected, long updated) {
+        throw new UnsupportedOperationException(); // TODO
+    }
+
+    @Override
+    public boolean compareAndSetRegister(String key, long expected, long updated) {
+        throw new UnsupportedOperationException(); // TODO
+    }
+
+    @Override
+    public long getRegister(String key) throws IOException {
+        throw new UnsupportedOperationException(); // TODO
+    }
 }

modules/repository-s3/src/main/java/org/elasticsearch/repositories/s3/S3BlobContainer.java

@@ -607,4 +607,19 @@ class S3BlobContainer extends AbstractBlobContainer {
             return Tuple.tuple(parts + 1, remaining);
         }
     }
+
+    @Override
+    public long compareAndExchangeRegister(String key, long expected, long updated) {
+        throw new UnsupportedOperationException(); // TODO
+    }
+
+    @Override
+    public boolean compareAndSetRegister(String key, long expected, long updated) {
+        throw new UnsupportedOperationException(); // TODO
+    }
+
+    @Override
+    public long getRegister(String key) throws IOException {
+        throw new UnsupportedOperationException(); // TODO
+    }
 }

modules/repository-url/src/main/java/org/elasticsearch/common/blobstore/url/URLBlobContainer.java

@@ -147,4 +147,9 @@ public class URLBlobContainer extends AbstractBlobContainer {
         }
     }
 
+    @Override
+    public long compareAndExchangeRegister(String key, long expected, long updated) {
+        throw new UnsupportedOperationException("URL repository doesn't support this operation");
+    }
+
 }

plugins/repository-hdfs/src/main/java/org/elasticsearch/repositories/hdfs/HdfsBlobContainer.java

@@ -316,4 +316,9 @@ final class HdfsBlobContainer extends AbstractBlobContainer {
             });
         }
     }
+
+    @Override
+    public long compareAndExchangeRegister(String key, long expected, long updated) {
+        throw new UnsupportedOperationException("HDFS repositories do not support this operation");
+    }
 }

server/src/main/java/org/elasticsearch/common/blobstore/BlobContainer.java

@@ -194,4 +194,40 @@ public interface BlobContainer {
      * @throws  IOException if there were any failures in reading from the blob container.
      */
     Map<String, BlobMetadata> listBlobsByPrefix(String blobNamePrefix) throws IOException;
+
+    /**
+     * Atomically sets the value stored at the given key to {@code updated} if the {@code current value == expected}.
+     * Keys not yet used start at initial value 0. Returns the current value (before it was updated).
+     *
+     * @param key key of the value to update
+     * @param expected the expected value
+     * @param updated the new value
+     * @return the value read from the register (before it was updated)
+     */
+    long compareAndExchangeRegister(String key, long expected, long updated) throws IOException;
+
+    /**
+     * Atomically sets the value stored at the given key to {@code updated} if the {@code current value == expected}.
+     * Keys not yet used start at initial value 0.
+     *
+     * @param key key of the value to update
+     * @param expected the expected value
+     * @param updated the new value
+     * @return true if successful, false if the expected value did not match the updated value
+     */
+    default boolean compareAndSetRegister(String key, long expected, long updated) throws IOException {
+        return compareAndExchangeRegister(key, expected, updated) == expected;
+    }
+
+    /**
+     * Gets the value set by {@link #compareAndSetRegister(String, long, long)} for a given key.
+     * If a key has not yet been used, the initial value is 0.
+     *
+     * @param key key of the value to get
+     * @return value found
+     */
+    default long getRegister(String key) throws IOException {
+        return compareAndExchangeRegister(key, 0, 0);
+    }
+
 }

server/src/main/java/org/elasticsearch/common/blobstore/fs/FsBlobContainer.java

@@ -20,15 +20,21 @@ import org.elasticsearch.common.blobstore.support.BlobMetadata;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.Streams;
+import org.elasticsearch.common.util.concurrent.KeyedLock;
 import org.elasticsearch.core.CheckedConsumer;
 import org.elasticsearch.core.IOUtils;
+import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Strings;
+import org.elasticsearch.core.SuppressForbidden;
 
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.nio.channels.Channels;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileLock;
 import java.nio.channels.SeekableByteChannel;
 import java.nio.file.AccessDeniedException;
 import java.nio.file.DirectoryStream;
@@ -369,4 +375,54 @@ public class FsBlobContainer extends AbstractBlobContainer {
     private static OutputStream blobOutputStream(Path file) throws IOException {
         return Files.newOutputStream(file, StandardOpenOption.CREATE_NEW);
     }
+
+    private static final KeyedLock<String> registerLocks = new KeyedLock<>();
+
+    @Override
+    @SuppressForbidden(reason = "write to channel that we have open for locking purposes already directly")
+    public long compareAndExchangeRegister(String key, long expected, long updated) throws IOException {
+        try (
+            FileChannel channel = openOrCreateAtomic(path.resolve(key));
+            FileLock ignored1 = channel.lock();
+            Releasable ignored2 = registerLocks.acquire(key)
+        ) {
+            final ByteBuffer buf = ByteBuffer.allocate(Long.BYTES);
+            final long found;
+            while (buf.remaining() > 0) {
+                if (channel.read(buf) == -1) {
+                    break;
+                }
+            }
+            if (buf.position() == 0) {
+                found = 0L;
+            } else if (buf.position() == Long.BYTES) {
+                found = buf.getLong(0);
+                buf.clear();
+                if (channel.read(buf) != -1) {
+                    throw new IllegalStateException("Read file of length greater than [" + Long.BYTES + "] for [" + key + "]");
+                }
+            } else {
+                throw new IllegalStateException("Read file of length [" + buf.position() + "] for [" + key + "]");
+            }
+            if (found == expected) {
+                buf.clear().putLong(updated).flip();
+                while (buf.remaining() > 0) {
+                    channel.write(buf, buf.position());
+                }
+                channel.force(true);
+            }
+            return found;
+        }
+    }
+
+    private static FileChannel openOrCreateAtomic(Path path) throws IOException {
+        try {
+            if (Files.exists(path) == false) {
+                return FileChannel.open(path, StandardOpenOption.READ, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE);
+            }
+        } catch (FileAlreadyExistsException e) {
+            // ok, created concurrently
+        }
+        return FileChannel.open(path, StandardOpenOption.READ, StandardOpenOption.WRITE);
+    }
 }

server/src/main/java/org/elasticsearch/common/blobstore/support/FilterBlobContainer.java

@@ -101,4 +101,19 @@ public abstract class FilterBlobContainer implements BlobContainer {
     public Map<String, BlobMetadata> listBlobsByPrefix(String blobNamePrefix) throws IOException {
         return delegate.listBlobsByPrefix(blobNamePrefix);
     }
+
+    @Override
+    public long compareAndExchangeRegister(String key, long expected, long updated) throws IOException {
+        return delegate.compareAndExchangeRegister(key, expected, updated);
+    }
+
+    @Override
+    public boolean compareAndSetRegister(String key, long expected, long updated) throws IOException {
+        return delegate.compareAndSetRegister(key, expected, updated);
+    }
+
+    @Override
+    public long getRegister(String key) throws IOException {
+        return delegate.getRegister(key);
+    }
 }

server/src/test/java/org/elasticsearch/common/blobstore/fs/FsBlobContainerTests.java

@@ -11,6 +11,7 @@ import org.apache.lucene.tests.mockfile.FilterFileSystemProvider;
 import org.apache.lucene.tests.mockfile.FilterSeekableByteChannel;
 import org.apache.lucene.tests.util.LuceneTestCase;
 import org.elasticsearch.common.blobstore.BlobPath;
+import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.io.Streams;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.PathUtils;
@@ -94,6 +95,48 @@ public class FsBlobContainerTests extends ESTestCase {
         assertThat(FsBlobContainer.isTempBlobName(tempBlobName), is(true));
     }
 
+    public void testCompareAndExchange() throws Exception {
+        final Path path = PathUtils.get(createTempDir().toString());
+        final FsBlobContainer container = new FsBlobContainer(
+            new FsBlobStore(randomIntBetween(1, 8) * 1024, path, false),
+            BlobPath.EMPTY,
+            path
+        );
+
+        final String key = randomAlphaOfLength(10);
+        final AtomicLong expectedValue = new AtomicLong();
+
+        for (int i = 0; i < 5; i++) {
+            switch (between(1, 4)) {
+                case 1 -> assertEquals(expectedValue.get(), container.getRegister(key));
+                case 2 -> assertFalse(
+                    container.compareAndSetRegister(key, randomValueOtherThan(expectedValue.get(), ESTestCase::randomLong), randomLong())
+                );
+                case 3 -> assertEquals(
+                    expectedValue.get(),
+                    container.compareAndExchangeRegister(
+                        key,
+                        randomValueOtherThan(expectedValue.get(), ESTestCase::randomLong),
+                        randomLong()
+                    )
+                );
+                case 4 -> {/* no-op */}
+            }
+
+            final var newValue = randomLong();
+            if (randomBoolean()) {
+                assertTrue(container.compareAndSetRegister(key, expectedValue.get(), newValue));
+            } else {
+                assertEquals(expectedValue.get(), container.compareAndExchangeRegister(key, expectedValue.get(), newValue));
+            }
+            expectedValue.set(newValue);
+        }
+
+        final byte[] corruptContents = new byte[9];
+        container.writeBlob(key, new BytesArray(corruptContents, 0, randomFrom(1, 7, 9)), false);
+        expectThrows(IllegalStateException.class, () -> container.compareAndExchangeRegister(key, expectedValue.get(), 0));
+    }
+
     static class MockFileSystemProvider extends FilterFileSystemProvider {
 
         final Consumer<Long> onRead;

x-pack/plugin/searchable-snapshots/src/test/java/org/elasticsearch/xpack/searchablesnapshots/cache/common/TestUtils.java

@@ -228,6 +228,11 @@ public final class TestUtils {
             throw unsupportedException();
         }
 
+        @Override
+        public long compareAndExchangeRegister(String key, long expected, long updated) throws IOException {
+            throw unsupportedException();
+        }
+
         private UnsupportedOperationException unsupportedException() {
             assert false : "this operation is not supported and should have not be called";
             return new UnsupportedOperationException("This operation is not supported");

x-pack/plugin/snapshot-repo-test-kit/src/internalClusterTest/java/org/elasticsearch/repositories/blobstore/testkit/RepositoryAnalysisFailureIT.java

@@ -537,6 +537,12 @@ public class RepositoryAnalysisFailureIT extends AbstractSnapshotIntegTestCase {
             blobMetadataByName.keySet().removeIf(s -> s.startsWith(blobNamePrefix) == false);
             return blobMetadataByName;
         }
+
+        @Override
+        public long compareAndExchangeRegister(String key, long expected, long updated) {
+            assert false : "should not have been called";
+            throw new UnsupportedOperationException();
+        }
     }
 
 }

x-pack/plugin/snapshot-repo-test-kit/src/internalClusterTest/java/org/elasticsearch/repositories/blobstore/testkit/RepositoryAnalysisSuccessIT.java

@@ -404,6 +404,12 @@ public class RepositoryAnalysisSuccessIT extends AbstractSnapshotIntegTestCase {
             blobMetadataByName.keySet().removeIf(s -> s.startsWith(blobNamePrefix) == false);
             return blobMetadataByName;
         }
+
+        @Override
+        public long compareAndExchangeRegister(String key, long expected, long updated) {
+            assert false : "should not have been called";
+            throw new UnsupportedOperationException();
+        }
     }
 
 }