WIP #2990: make LucenePartitioner thread safe

This changes the concurrentUpdate to run with a single partition,
and adds usages of AsyncLock to ensure that updates to the partition
metadata work correctly.
Before calling the issue complete an additional test should be added
that does concurrent inserts, and one for concurrent deletes.

Author: ScottDugas

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordCoreInternalException.java

@@ -0,0 +1,38 @@
+/*
+ * RecordCoreInternalException.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.record;
+
+import com.apple.foundationdb.annotation.API;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Exception thrown when an inconsistency in core record layer behavior is detected.
+ */
+@API(API.Status.STABLE)
+@SuppressWarnings("serial")
+public class RecordCoreInternalException extends RecordCoreException {
+
+    public RecordCoreInternalException(@Nonnull final String msg, @Nullable final Object... keyValues) {
+        super(msg, keyValues);
+    }
+}

fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneIndexMaintainer.java

@@ -246,6 +246,14 @@ private void insertField(LuceneDocumentFromRecord.DocumentField field, final Doc
         }
     }
 
+    <M extends Message> void writeDocument(final FDBIndexableRecord<M> newRecord, final Map.Entry<Tuple, List<LuceneDocumentFromRecord.DocumentField>> entry, final Integer partitionId) {
+        try {
+            writeDocument(entry.getValue(), entry.getKey(), partitionId, newRecord.getPrimaryKey());
+        } catch (IOException e) {
+            throw LuceneExceptions.toRecordCoreException("Issue updating new index keys", e, "newRecord", newRecord.getPrimaryKey());
+        }
+    }
+
     @SuppressWarnings("PMD.CloseResource")
     private void writeDocument(@Nonnull List<LuceneDocumentFromRecord.DocumentField> fields,
                                Tuple groupingKey,
@@ -484,14 +492,11 @@ <M extends Message> CompletableFuture<Void> update(@Nullable FDBIndexableRecord<
                 AsyncUtil.whenAll(newRecordFields.entrySet().stream().map(entry -> {
                     try {
                         return tryDeleteInWriteOnlyMode(Objects.requireNonNull(newRecord), entry.getKey()).thenCompose(countDeleted ->
-                                partitioner.addToAndSavePartitionMetadata(newRecord, entry.getKey(), destinationPartitionIdHint).thenApply(partitionId -> {
-                                    try {
-                                        writeDocument(entry.getValue(), entry.getKey(), partitionId, newRecord.getPrimaryKey());
-                                    } catch (IOException e) {
-                                        throw LuceneExceptions.toRecordCoreException("Issue updating new index keys", e, "newRecord", newRecord.getPrimaryKey());
-                                    }
-                                    return null;
-                                }));
+                                partitioner.addToAndSavePartitionMetadata(newRecord, entry.getKey(), destinationPartitionIdHint)
+                                        .thenApply(partitionId -> {
+                                            writeDocument(newRecord, entry, partitionId);
+                                            return null;
+                                        }));
                     } catch (IOException e) {
                         throw LuceneExceptions.toRecordCoreException("Issue updating", e, "record", Objects.requireNonNull(newRecord).getPrimaryKey());
                     }
@@ -538,19 +543,22 @@ private <M extends Message> CompletableFuture<Integer> tryDelete(@Nonnull FDBInd
         }
 
         // partitioned
-        return partitioner.tryGetPartitionInfo(record, groupingKey).thenApply(partitionInfo -> {
+        return partitioner.tryGetPartitionInfo(record, groupingKey).thenCompose(partitionInfo -> {
+            // this might be 0 when in writeOnly mode, but otherwise should not happen.
             if (partitionInfo != null) {
                 try {
                     int countDeleted = deleteDocument(groupingKey, partitionInfo.getId(), record.getPrimaryKey());
                     if (countDeleted > 0) {
-                        partitioner.decrementCountAndSave(groupingKey, partitionInfo, countDeleted);
+                        return partitioner.decrementCountAndSave(groupingKey, countDeleted, partitionInfo.getId())
+                                .thenApply(vignore -> countDeleted);
+                    } else {
+                        return CompletableFuture.completedFuture(countDeleted);
                     }
-                    return countDeleted;
                 } catch (IOException e) {
                     throw LuceneExceptions.toRecordCoreException("Issue deleting", e, "record", record.getPrimaryKey());
                 }
             }
-            return 0;
+            return CompletableFuture.completedFuture(0);
         });
     }
 

fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LucenePartitioner.java

@@ -33,13 +33,15 @@
 import com.apple.foundationdb.record.PipelineOperation;
 import com.apple.foundationdb.record.RecordCoreArgumentException;
 import com.apple.foundationdb.record.RecordCoreException;
+import com.apple.foundationdb.record.RecordCoreInternalException;
 import com.apple.foundationdb.record.RecordCursor;
 import com.apple.foundationdb.record.RecordCursorContinuation;
 import com.apple.foundationdb.record.RecordCursorEndContinuation;
 import com.apple.foundationdb.record.RecordCursorStartContinuation;
 import com.apple.foundationdb.record.ScanProperties;
 import com.apple.foundationdb.record.TupleRange;
 import com.apple.foundationdb.record.cursors.ChainedCursor;
+import com.apple.foundationdb.record.locking.LockIdentifier;
 import com.apple.foundationdb.record.logging.KeyValueLogMessage;
 import com.apple.foundationdb.record.logging.LogMessageKeys;
 import com.apple.foundationdb.record.lucene.directory.FDBDirectoryManager;
@@ -81,7 +83,6 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Comparator;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
@@ -489,28 +490,30 @@ public <M extends Message> CompletableFuture<Integer> addToAndSavePartitionMetad
     private CompletableFuture<Integer> addToAndSavePartitionMetadata(@Nonnull final Tuple groupingKey,
                                                                      @Nonnull final Tuple partitioningKey,
                                                                      @Nullable final Integer assignedPartitionIdOverride) {
-
-        final CompletableFuture<LucenePartitionInfoProto.LucenePartitionInfo> assignmentFuture;
-        if (assignedPartitionIdOverride != null) {
-            assignmentFuture = getPartitionMetaInfoById(assignedPartitionIdOverride, groupingKey);
-        } else {
-            assignmentFuture = getOrCreatePartitionInfo(groupingKey, partitioningKey);
-        }
-        return assignmentFuture.thenApply(assignedPartition -> {
-            // assignedPartition is not null, since a new one is created by the previous call if none exist
-            LucenePartitionInfoProto.LucenePartitionInfo.Builder builder = Objects.requireNonNull(assignedPartition).toBuilder();
-            builder.setCount(assignedPartition.getCount() + 1);
-            if (isOlderThan(partitioningKey, assignedPartition)) {
-                // clear the previous key
-                state.context.ensureActive().clear(partitionMetadataKeyFromPartitioningValue(groupingKey, getPartitionKey(assignedPartition)));
-                builder.setFrom(ByteString.copyFrom(partitioningKey.pack()));
-            }
-            if (isNewerThan(partitioningKey, assignedPartition)) {
-                builder.setTo(ByteString.copyFrom(partitioningKey.pack()));
-            }
-            savePartitionMetadata(groupingKey, builder);
-            return assignedPartition.getId();
-        });
+        return state.context.doWithWriteLock(new LockIdentifier(partitionMetadataSubspace(groupingKey)),
+                () -> {
+                    final CompletableFuture<LucenePartitionInfoProto.LucenePartitionInfo> assignmentFuture;
+                    if (assignedPartitionIdOverride != null) {
+                        assignmentFuture = getPartitionMetaInfoById(assignedPartitionIdOverride, groupingKey);
+                    } else {
+                        assignmentFuture = getOrCreatePartitionInfo(groupingKey, partitioningKey);
+                    }
+                    return assignmentFuture.thenApply(assignedPartition -> {
+                        // assignedPartition is not null, since a new one is created by the previous call if none exist
+                        LucenePartitionInfoProto.LucenePartitionInfo.Builder builder = Objects.requireNonNull(assignedPartition).toBuilder();
+                        builder.setCount(assignedPartition.getCount() + 1);
+                        if (isOlderThan(partitioningKey, assignedPartition)) {
+                            // clear the previous key
+                            state.context.ensureActive().clear(partitionMetadataKeyFromPartitioningValue(groupingKey, getPartitionKey(assignedPartition)));
+                            builder.setFrom(ByteString.copyFrom(partitioningKey.pack()));
+                        }
+                        if (isNewerThan(partitioningKey, assignedPartition)) {
+                            builder.setTo(ByteString.copyFrom(partitioningKey.pack()));
+                        }
+                        savePartitionMetadata(groupingKey, builder);
+                        return assignedPartition.getId();
+                    });
+                });
     }
 
     /**
@@ -525,6 +528,10 @@ byte[] partitionMetadataKeyFromPartitioningValue(@Nonnull Tuple groupKey, @Nonnu
         return state.indexSubspace.pack(partitionMetadataKeyTuple(groupKey, partitionKey));
     }
 
+    Subspace partitionMetadataSubspace(@Nonnull Tuple groupKey) {
+        return state.indexSubspace.subspace(groupKey.add(PARTITION_META_SUBSPACE));
+    }
+
     private static Tuple partitionMetadataKeyTuple(final @Nonnull Tuple groupKey, @Nonnull Tuple partitionKey) {
         return groupKey.add(PARTITION_META_SUBSPACE).addAll(partitionKey);
     }
@@ -600,22 +607,30 @@ <M extends Message> CompletableFuture<LucenePartitionInfoProto.LucenePartitionIn
      * decrement the doc count of a partition, and save its partition metadata.
      *
      * @param groupingKey grouping key
-     * @param partitionInfo partition metadata
      * @param amount amount to subtract from the doc count
+     * @param partitionId the id of the partition to decrement
      */
-    void decrementCountAndSave(@Nonnull Tuple groupingKey,
-                               @Nonnull LucenePartitionInfoProto.LucenePartitionInfo partitionInfo,
-                               int amount) {
-        LucenePartitionInfoProto.LucenePartitionInfo.Builder builder = Objects.requireNonNull(partitionInfo).toBuilder();
-        // note that the to/from of the partition do not get updated, since that would require us to know what the next potential boundary
-        // value(s) are. The values, nonetheless, remain valid.
-        builder.setCount(partitionInfo.getCount() - amount);
-
-        if (builder.getCount() < 0) {
-            // should never happen
-            throw new RecordCoreException("Issue updating Lucene partition metadata (resulting count < 0)", LogMessageKeys.PARTITION_ID, partitionInfo.getId());
-        }
-        savePartitionMetadata(groupingKey, builder);
+    CompletableFuture<Void> decrementCountAndSave(@Nonnull Tuple groupingKey,
+                                                  int amount, final int partitionId) {
+        return state.context.doWithWriteLock(new LockIdentifier(partitionMetadataSubspace(groupingKey)),
+                () -> getPartitionMetaInfoById(partitionId, groupingKey).thenAccept(serialized -> {
+                    if (serialized == null) {
+                        throw new RecordCoreInternalException("Lucene partition metadata changed during delete")
+                                .addLogInfo(LogMessageKeys.INDEX_NAME, state.index.getName())
+                                .addLogInfo(LogMessageKeys.INDEX_SUBSPACE, state.indexSubspace);
+                    }
+                    LucenePartitionInfoProto.LucenePartitionInfo.Builder builder = Objects.requireNonNull(serialized).toBuilder();
+                    // note that the to/from of the partition do not get updated, since that would require us to know
+                    // what the next potential boundary value(s) are. The values, nonetheless, remain valid.
+                    builder.setCount(serialized.getCount() - amount);
+
+                    if (builder.getCount() < 0) {
+                        // should never happen
+                        throw new RecordCoreInternalException("Issue updating Lucene partition metadata (resulting count < 0)",
+                                LogMessageKeys.PARTITION_ID, partitionId);
+                    }
+                    savePartitionMetadata(groupingKey, builder);
+                }));
     }
 
     /**
@@ -1024,7 +1039,7 @@ private CompletableFuture<Integer> moveDocsFromPartition(@Nonnull final LuceneRe
 
         timings.initializationNanos = System.nanoTime();
         fetchedRecordsFuture = fetchedRecordsFuture.whenComplete((ignored, throwable) -> cursor.close());
-        return fetchedRecordsFuture.thenCompose(records -> {
+        return fetchedRecordsFuture.thenApply(records -> {
             timings.searchNanos = System.nanoTime();
             if (records.size() == 0) {
                 throw new RecordCoreException("Unexpected error: 0 records fetched. repartitionContext {}", repartitioningContext);
@@ -1041,7 +1056,7 @@ private CompletableFuture<Integer> moveDocsFromPartition(@Nonnull final LuceneRe
                 if (LOGGER.isDebugEnabled()) {
                     LOGGER.debug("no records to move, partition {}", partitionInfo);
                 }
-                return CompletableFuture.completedFuture(0);
+                return 0;
             }
 
             // reset partition info
@@ -1092,40 +1107,43 @@ private CompletableFuture<Integer> moveDocsFromPartition(@Nonnull final LuceneRe
             }
             long updateStart = System.nanoTime();
 
-            Iterator<? extends FDBIndexableRecord<Message>> recordIterator = records.iterator();
             final int destinationPartitionId = destinationPartition.getId();
-            return AsyncUtil.whileTrue(() -> indexMaintainer.update(null, recordIterator.next(), destinationPartitionId)
-                    .thenApply(ignored -> recordIterator.hasNext()), state.context.getExecutor())
-                    .thenApply(ignored -> {
-                        if (LOGGER.isDebugEnabled()) {
-                            long updateNanos = System.nanoTime();
-                            final KeyValueLogMessage logMessage = repartitionLogMessage("Repartitioned records", groupingKey, records.size(), partitionInfo);
-                            logMessage.addKeyAndValue("totalMicros", TimeUnit.NANOSECONDS.toMicros(updateNanos - timings.startNanos));
-                            logMessage.addKeyAndValue("initializationMicros", TimeUnit.NANOSECONDS.toMicros(timings.initializationNanos - timings.startNanos));
-                            logMessage.addKeyAndValue("searchMicros", TimeUnit.NANOSECONDS.toMicros(timings.searchNanos - timings.initializationNanos));
-                            logMessage.addKeyAndValue("clearInfoMicros", TimeUnit.NANOSECONDS.toMicros(timings.clearInfoNanos - timings.searchNanos));
-                            if (timings.emptyingNanos > 0) {
-                                logMessage.addKeyAndValue("emptyingMicros", TimeUnit.NANOSECONDS.toMicros(timings.emptyingNanos - timings.clearInfoNanos));
-                            }
-                            if (timings.deleteNanos > 0) {
-                                logMessage.addKeyAndValue("deleteMicros", TimeUnit.NANOSECONDS.toMicros(timings.deleteNanos - timings.clearInfoNanos));
-                            }
-                            if (timings.metadataUpdateNanos > 0) {
-                                logMessage.addKeyAndValue("metadataUpdateMicros", TimeUnit.NANOSECONDS.toMicros(timings.metadataUpdateNanos - timings.deleteNanos));
-                            }
-                            if (timings.createPartitionNanos > 0) {
-                                logMessage.addKeyAndValue("createPartitionMicros", TimeUnit.NANOSECONDS.toMicros(timings.createPartitionNanos - endCleanupNanos));
-                            }
-                            logMessage.addKeyAndValue("updateMicros", TimeUnit.NANOSECONDS.toMicros(updateNanos - updateStart));
-                            if (timerSnapshot != null && state.context.getTimer() != null) {
-                                logMessage.addKeysAndValues(
-                                        StoreTimer.getDifference(state.context.getTimer(), timerSnapshot)
-                                                .getKeysAndValues());
-                            }
-                            LOGGER.debug(logMessage.toString());
-                        }
-                        return records.size();
-                    });
+            for (FDBIndexableRecord<Message> record : records) {
+                LuceneDocumentFromRecord.getRecordFields(state.index.getRootExpression(), record)
+                        .entrySet().forEach(entry -> {
+                            indexMaintainer.writeDocument(record, entry, destinationPartitionId);
+                            // TODO could update the partition once
+                            addToAndSavePartitionMetadata(record, groupingKey, destinationPartitionId);
+                        });
+            }
+            if (LOGGER.isDebugEnabled()) {
+                long updateNanos = System.nanoTime();
+                final KeyValueLogMessage logMessage = repartitionLogMessage("Repartitioned records", groupingKey, records.size(), partitionInfo);
+                logMessage.addKeyAndValue("totalMicros", TimeUnit.NANOSECONDS.toMicros(updateNanos - timings.startNanos));
+                logMessage.addKeyAndValue("initializationMicros", TimeUnit.NANOSECONDS.toMicros(timings.initializationNanos - timings.startNanos));
+                logMessage.addKeyAndValue("searchMicros", TimeUnit.NANOSECONDS.toMicros(timings.searchNanos - timings.initializationNanos));
+                logMessage.addKeyAndValue("clearInfoMicros", TimeUnit.NANOSECONDS.toMicros(timings.clearInfoNanos - timings.searchNanos));
+                if (timings.emptyingNanos > 0) {
+                    logMessage.addKeyAndValue("emptyingMicros", TimeUnit.NANOSECONDS.toMicros(timings.emptyingNanos - timings.clearInfoNanos));
+                }
+                if (timings.deleteNanos > 0) {
+                    logMessage.addKeyAndValue("deleteMicros", TimeUnit.NANOSECONDS.toMicros(timings.deleteNanos - timings.clearInfoNanos));
+                }
+                if (timings.metadataUpdateNanos > 0) {
+                    logMessage.addKeyAndValue("metadataUpdateMicros", TimeUnit.NANOSECONDS.toMicros(timings.metadataUpdateNanos - timings.deleteNanos));
+                }
+                if (timings.createPartitionNanos > 0) {
+                    logMessage.addKeyAndValue("createPartitionMicros", TimeUnit.NANOSECONDS.toMicros(timings.createPartitionNanos - endCleanupNanos));
+                }
+                logMessage.addKeyAndValue("updateMicros", TimeUnit.NANOSECONDS.toMicros(updateNanos - updateStart));
+                if (timerSnapshot != null && state.context.getTimer() != null) {
+                    logMessage.addKeysAndValues(
+                            StoreTimer.getDifference(state.context.getTimer(), timerSnapshot)
+                                    .getKeysAndValues());
+                }
+                LOGGER.debug(logMessage.toString());
+            }
+            return records.size();
         });
     }