Test attempting to reproduce FDB issue #11500

Issue: apple/foundationdb#11500

fdb-extensions/src/test/java/com/apple/foundationdb/LoadUnderBatchGrvTest.java

@@ -0,0 +1,196 @@
+/*
+ * LoadUnderBatchGrvTest.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2015-2024 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed 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 com.apple.foundationdb;
+
+import com.apple.foundationdb.async.AsyncIterator;
+import com.apple.foundationdb.subspace.Subspace;
+import com.apple.foundationdb.system.SystemKeyspace;
+import com.apple.foundationdb.test.TestDatabaseExtension;
+import com.apple.foundationdb.test.TestSubspaceExtension;
+import com.apple.foundationdb.tuple.Tuple;
+import com.apple.test.SuperSlow;
+import com.apple.test.Tags;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import java.util.HashSet;
+import java.util.UUID;
+import java.util.Vector;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Test that FDB behaves well when saturated under batch priority.
+ */
+@Tag(Tags.RequiresFDB)
+public class LoadUnderBatchGrvTest {
+
+    private static final Logger logger = LogManager.getLogger(LoadUnderBatchGrvTest.class);
+    @RegisterExtension
+    static final TestDatabaseExtension dbExtension = new TestDatabaseExtension();
+    @RegisterExtension
+    TestSubspaceExtension rsSubspaceExtension = new TestSubspaceExtension(dbExtension);
+
+    enum BatchActivity {
+        RegularKey,
+        RegularKeyAtSnapshot,
+        RegularRange,
+        RegularRangeAtSnapshot,
+        MetadataVersionKey,
+        MetadataVersionKeyAtSnapshot,
+    }
+
+    @Test
+    @SuperSlow
+    void loadTest() {
+        AtomicBoolean stopped = new AtomicBoolean();
+        final Database database = dbExtension.getDatabase();
+        final Subspace subspace = rsSubspaceExtension.getSubspace();
+        final byte[] value = new byte[1000];
+        ThreadLocalRandom.current().nextBytes(value);
+        final ScheduledExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(2);
+        final int cachedKeyCount = 1000;
+        final Vector<UUID> keys = new Vector<>(cachedKeyCount);
+        for (int i = 0; i < cachedKeyCount; i++) {
+            keys.add(UUID.randomUUID());
+        }
+        final ForkJoinPool writePool = new ForkJoinPool(100);
+        final ForkJoinPool readPool = new ForkJoinPool(BatchActivity.values().length * 5);
+        CompletableFuture<Void> start = new CompletableFuture<>();
+        CompletableFuture<Void> done = new CompletableFuture<>();
+        // generate batch grv load
+        for (final BatchActivity batchActivity : BatchActivity.values()) {
+            CompletableFuture.runAsync(() -> {
+                while (!stopped.get()) {
+                    //try {
+                    try (Transaction transaction = database.createTransaction()) {
+                        transaction.options().setPriorityBatch();
+                        logger.info(batchActivity + " starting");
+                        doGet(batchActivity, transaction, subspace, keys)
+//                        database.runAsync(transaction -> doGet(batchActivity, transaction, subspace, keys))
+                                .get(40, TimeUnit.SECONDS);
+                        Thread.sleep(10);
+                        logger.debug("Got " + batchActivity);
+                    } catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        stopped.set(true);
+                        throw new RuntimeException("Interrupted");
+                    } catch (ExecutionException e) {
+                        done.complete(null);
+                        logger.info(batchActivity + " failed", e);
+                        throw new RuntimeException(e);
+                    } catch (TimeoutException e) {
+                        done.complete(null);
+                        logger.info(batchActivity + " Timed out", e);
+                        throw new RuntimeException(e);
+                    }
+                }
+            }, readPool);
+        }
+
+        // generate a bunch of load at default priority
+        final ScheduledFuture<?> scheduledFuture = scheduledExecutorService.scheduleAtFixedRate(() ->
+                CompletableFuture.runAsync(() ->
+                        generateLoad(stopped, database, subspace, value, keys, start), writePool),
+                0, 2000, TimeUnit.MILLISECONDS);
+        done.join();
+        stopped.set(true);
+        scheduledFuture.cancel(true);
+    }
+
+    private static void generateLoad(final AtomicBoolean stopped, final Database database, final Subspace subspace, final byte[] value, final Vector<UUID> keys, final CompletableFuture<Void> start) {
+        final UUID worker = UUID.randomUUID();
+        logger.debug("Starting " + worker);
+        while (!stopped.get()) {
+            final HashSet<UUID> localKeys = new HashSet<>();
+            database.runAsync(transaction -> {
+                localKeys.clear(); // if the transaction failed and was retried by runAsync
+                for (int i = 0; i < 15; i++) {
+                    final UUID key = UUID.randomUUID();
+                    localKeys.add(key);
+                    transaction.set(subspace.pack(Tuple.from(key)), value);
+                }
+                for (int i = 0; i < 10; i++) {
+                    transaction.get(subspace.pack(Tuple.from(randomFrom(keys))));
+                }
+                return null;
+            }).whenComplete((vignore, err) -> {
+                if (err != null) {
+                    localKeys.stream().limit(3)
+                            .forEach(key -> setRandom(key, keys));
+                    start.complete(null);
+                }
+                //logger.debug("Committed " + worker);
+            });
+        }
+    }
+
+    private static CompletableFuture<byte[]> doGet(final BatchActivity batchActivity, final Transaction transaction,
+                                                   final Subspace subspace, final Vector<UUID> keys) {
+        switch (batchActivity) {
+            case RegularKey:
+                return transaction.get(subspace.pack(Tuple.from(randomFrom(keys))));
+            case RegularKeyAtSnapshot:
+                return transaction.snapshot().get(subspace.pack(Tuple.from(randomFrom(keys))));
+            case RegularRange:
+                return firstValueOfRandomRange(transaction, subspace, keys);
+            case RegularRangeAtSnapshot:
+                return firstValueOfRandomRange(transaction.snapshot(), subspace, keys);
+            case MetadataVersionKey:
+                return transaction.get(SystemKeyspace.METADATA_VERSION_KEY);
+            case MetadataVersionKeyAtSnapshot:
+                return transaction.snapshot().get(SystemKeyspace.METADATA_VERSION_KEY);
+            default:
+                throw new RuntimeException("Unexpected batch activity " + batchActivity);
+        }
+    }
+
+    private static CompletableFuture<byte[]> firstValueOfRandomRange(final ReadTransaction transaction, final Subspace subspace, final Vector<UUID> keys) {
+        Tuple key1 = Tuple.from(randomFrom(keys));
+        Tuple key2 = Tuple.from(randomFrom(keys));
+        AsyncIterator<KeyValue> iterator;
+        if (key1.compareTo(key2) > 0) {
+            iterator = transaction.getRange(subspace.pack(key2), subspace.pack(key1)).iterator();
+        } else {
+            iterator = transaction.getRange(subspace.pack(key1), subspace.pack(key2)).iterator();
+        }
+        return iterator.onHasNext().thenApply(hasNext -> hasNext ? iterator.next().getValue() : null);
+    }
+
+    private static UUID randomFrom(final Vector<UUID> keys) {
+        return keys.get(ThreadLocalRandom.current().nextInt(keys.size()));
+    }
+
+    private static UUID setRandom(final UUID key, final Vector<UUID> keys) {
+        return keys.set(ThreadLocalRandom.current().nextInt(keys.size()), key);
+    }
+}