Add RecordCountKey State (#3334)

This adds a mechanism for controlling the state of the _deprecated_
RecordCountKey, that can be set on the RecordMetaData.

Note: This is different from `COUNT` indexes.

Resolves #3326 

As noted in the javadocs:

The new state is controlled by
`FDBRecordStore.updateRecordCountStateAsync`, and allows:
1. Changing a READABLE RecordCount to WRITE_ONLY
2. Changing a WRITE_ONLY RecordCount back to READABLE
3. Changing any state to DISABLED

The default state is READABLE, and it will stay READABLE unless
explicitly updated.
If the RecordCount is WRITE_ONLY then it will maintain the count, but
not allow querying it.
If the RecordCount is DISABLED it will clear the data, and stop
maintaining the values.


In order to support this new state, a new `FormatVersion` has been
added: `RECORD_COUNT_STATE`
In order to use the new functionality, the stores must be built with at
least this format version.

---------

Co-authored-by: Alec Grieser <alloc@apple.com>

Author: ScottDugas

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

@@ -559,13 +559,21 @@ protected <M extends Message> CompletableFuture<FDBStoredRecord<M>> saveTypedRec
 
     @SuppressWarnings("PMD.CloseResource")
     private <M extends Message> void addRecordCount(@Nonnull RecordMetaData metaData, @Nonnull FDBStoredRecord<M> rec, @Nonnull byte[] increment) {
-        if (metaData.getRecordCountKey() == null) {
-            return;
+        if (metaData.getRecordCountKey() != null) {
+            beginRecordStoreStateRead();
+            try {
+                RecordMetaDataProto.DataStoreInfo header = recordStoreStateRef.get().getStoreHeader();
+                // We do not need to check the format version here. In order for it to be DISABLED we would have to be
+                // on a format version that supports such a state.
+                if (header.getRecordCountState() != RecordMetaDataProto.DataStoreInfo.RecordCountState.DISABLED) {
+                    Key.Evaluated subkey = metaData.getRecordCountKey().evaluateSingleton(rec);
+                    final byte[] keyBytes = getSubspace().pack(Tuple.from(RECORD_COUNT_KEY).addAll(subkey.toTupleAppropriateList()));
+                    ensureContextActive().mutate(MutationType.ADD, keyBytes, increment);
+                }
+            } finally {
+                endRecordStoreStateRead();
+            }
         }
-        final Transaction tr = ensureContextActive();
-        Key.Evaluated subkey = metaData.getRecordCountKey().evaluateSingleton(rec);
-        final byte[] keyBytes = getSubspace().pack(Tuple.from(RECORD_COUNT_KEY).addAll(subkey.toTupleAppropriateList()));
-        tr.mutate(MutationType.ADD, keyBytes, increment);
     }
 
     @Nullable
@@ -1851,7 +1859,9 @@ private Key.Evaluated deleteRecordsWhereCheckRecordTypes() {
             }
 
             final KeyExpression recordCountKey = getRecordMetaData().getRecordCountKey();
-            if (recordCountKey != null) {
+            if (recordCountKey != null
+                    // we don't need to call beginRecordStoreStateRead(), that is checked in deleteRecordsWhereAsync
+                    && recordStoreStateRef.get().getStoreHeader().getRecordCountState() != RecordMetaDataProto.DataStoreInfo.RecordCountState.DISABLED) {
                 final QueryToKeyMatcher.Match match = matcher.matchesSatisfyingQuery(recordCountKey);
                 if (match.getType() != QueryToKeyMatcher.MatchType.EQUALITY) {
                     throw new Query.InvalidExpressionException("Record count key not matching for deleteRecordsWhere");
@@ -2077,8 +2087,11 @@ private CompletableFuture<Void> run() {
                 context.clear(versionRange);
             }
 
+
             final KeyExpression recordCountKey = getRecordMetaData().getRecordCountKey();
-            if (recordCountKey != null) {
+            if (recordCountKey != null
+                    // we don't need to call beginRecordStoreStateRead(), that is checked in deleteRecordsWhereAsync
+                    && recordStoreStateRef.get().getStoreHeader().getRecordCountState() != RecordMetaDataProto.DataStoreInfo.RecordCountState.DISABLED) {
                 if (prefix.size() == recordCountKey.getColumnSize()) {
                     // Delete a single record used for counting
                     context.clear(getSubspace().pack(Tuple.from(RECORD_COUNT_KEY).addAll(prefix)));
@@ -2153,17 +2166,39 @@ private CompletableFuture<Long> estimateSize(@Nonnull Range range, long startTim
     @Override
     public CompletableFuture<Long> getSnapshotRecordCount(@Nonnull KeyExpression key, @Nonnull Key.Evaluated value,
                                                           @Nonnull IndexQueryabilityFilter indexQueryabilityFilter) {
-        if (getRecordMetaData().getRecordCountKey() != null) {
-            if (key.getColumnSize() != value.size()) {
-                throw recordCoreException("key and value are not the same size");
-            }
-            final ReadTransaction tr = context.readTransaction(true);
-            final Tuple subkey = Tuple.from(RECORD_COUNT_KEY).addAll(value.toTupleAppropriateList());
-            if (getRecordMetaData().getRecordCountKey().equals(key)) {
-                return tr.get(getSubspace().pack(subkey)).thenApply(FDBRecordStore::decodeRecordCount);
-            } else if (key.isPrefixKey(getRecordMetaData().getRecordCountKey())) {
-                AsyncIterable<KeyValue> kvs = tr.getRange(getSubspace().range(Tuple.from(RECORD_COUNT_KEY)));
-                return MoreAsyncUtil.reduce(getExecutor(), kvs.iterator(), 0L, (count, kv) -> count + decodeRecordCount(kv.getValue()));
+        final RecordMetaData recordMetaData = getRecordMetaData();
+        if (recordMetaData.getRecordCountKey() != null) {
+            beginRecordStoreStateRead();
+            boolean futureCreated = false;
+            try {
+                RecordMetaDataProto.DataStoreInfo header = recordStoreStateRef.get().getStoreHeader();
+                // We can always check the state, even if the formatVersion is older, because older versions will always
+                // have the default of READABLE
+                if (header.getRecordCountState() == RecordMetaDataProto.DataStoreInfo.RecordCountState.READABLE) {
+                    if (key.getColumnSize() != value.size()) {
+                        throw recordCoreException("key and value are not the same size");
+                    }
+                    final ReadTransaction tr = context.readTransaction(true);
+                    final Tuple subkey = Tuple.from(RECORD_COUNT_KEY).addAll(value.toTupleAppropriateList());
+                    if (recordMetaData.getRecordCountKey().equals(key)) {
+                        final CompletableFuture<Long> result = tr.get(getSubspace().pack(subkey))
+                                .thenApply(FDBRecordStore::decodeRecordCount)
+                                .whenComplete((ignored, error) -> endRecordStoreStateRead());
+                        futureCreated = true;
+                        return result;
+                    } else if (key.isPrefixKey(recordMetaData.getRecordCountKey())) {
+                        AsyncIterable<KeyValue> kvs = tr.getRange(getSubspace().range(Tuple.from(RECORD_COUNT_KEY)));
+                        final CompletableFuture<Long> result = MoreAsyncUtil.reduce(getExecutor(), kvs.iterator(), 0L,
+                                (count, kv) -> count + decodeRecordCount(kv.getValue()))
+                                .whenComplete((ignored, error) -> endRecordStoreStateRead());
+                        futureCreated = true;
+                        return result;
+                    }
+                }
+            } finally {
+                if (!futureCreated) {
+                    endRecordStoreStateRead();
+                }
             }
         }
         return evaluateAggregateFunction(Collections.emptyList(), IndexFunctionHelper.count(key),
@@ -3237,6 +3272,48 @@ public void clearHeaderUserField(@Nonnull String userField) {
         context.asyncToSync(FDBStoreTimer.Waits.WAIT_EDIT_HEADER_USER_FIELD, clearHeaderUserFieldAsync(userField));
     }
 
+    /**
+     * Update the RecordCount to have a new state.
+     * <p>
+     *     The state can go from {@code READABLE} to {@code WRITE_ONLY}, which is mostly useful to validate that if you
+     *     drop the {@link RecordMetaData#getRecordCountKey() recordCountKey} from the metadata, it won't have adverse
+     *     affects.
+     *     The state can always be set to {@code DISABLED}, at which point there is no way to any other state. This
+     *     limitation exists, in large part, because there is not a way to rebuild the count across indexes, and since
+     *     it is long deprecated, investing in that doesn't make sense.
+     * </p>
+     * @param newState the new state to update it to.
+     * @return a future once that completes once the state is updated, and the data is cleared.
+     */
+    public CompletableFuture<Void> updateRecordCountStateAsync(@Nonnull RecordMetaDataProto.DataStoreInfo.RecordCountState newState) {
+        return updateStoreHeaderAsync(builder -> {
+            if (!getFormatVersionEnum().isAtLeast(FormatVersion.RECORD_COUNT_STATE)) {
+                throw new RecordCoreException("Store does not support updating record count state")
+                        .addLogInfo(LogMessageKeys.FORMAT_VERSION, getFormatVersionEnum());
+            }
+            final RecordMetaDataProto.DataStoreInfo.RecordCountState existing = getRecordStoreState().getStoreHeader().getRecordCountState();
+            if (existing == newState) {
+                return builder;
+            }
+            boolean toWriteOnly = existing == RecordMetaDataProto.DataStoreInfo.RecordCountState.READABLE &&
+                    newState == RecordMetaDataProto.DataStoreInfo.RecordCountState.WRITE_ONLY;
+            boolean toReadable = existing == RecordMetaDataProto.DataStoreInfo.RecordCountState.WRITE_ONLY &&
+                    newState == RecordMetaDataProto.DataStoreInfo.RecordCountState.READABLE;
+            if (toWriteOnly || toReadable || newState == RecordMetaDataProto.DataStoreInfo.RecordCountState.DISABLED) {
+                builder.setRecordCountState(newState);
+                if (newState == RecordMetaDataProto.DataStoreInfo.RecordCountState.DISABLED) {
+                    // Note: getSubspace().range(tuple) does not include getSubspace().pack(tuple), but this does.
+                    // Ungrouped recordCountKey will be stored directly at getSubspace().pack(tuple).
+                    ensureContextActive().clear(Range.startsWith(getSubspace().pack(Tuple.from(RECORD_COUNT_KEY))));
+                }
+                return builder;
+            }
+            throw new RecordCoreException("Invalid state transition for RecordCountState")
+                    .addLogInfo(LogMessageKeys.OLD, existing)
+                    .addLogInfo(LogMessageKeys.NEW, newState);
+        });
+    }
+
     // Actually (1) writes the index state to the database and (2) updates the cached state with the new state
     @SuppressWarnings("PMD.CloseResource")
     private void updateIndexState(@Nonnull String indexName, byte[] indexKey, @Nonnull IndexState indexState) {
@@ -4841,6 +4918,8 @@ protected boolean checkPossiblyRebuildRecordCounts(@Nonnull RecordMetaData metaD
 
         if (rebuildRecordCounts) {
             // We want to clear all record counts.
+            // This code will leave data behind if the previous RecordCountKey was not grouped
+            // https://github.com/FoundationDB/fdb-record-layer/issues/3335
             if (existingStore) {
                 context.clear(getSubspace().range(Tuple.from(RECORD_COUNT_KEY)));
             }
@@ -4865,7 +4944,8 @@ protected boolean checkPossiblyRebuildRecordCounts(@Nonnull RecordMetaData metaD
     @SuppressWarnings("PMD.CloseResource")
     public void addRebuildRecordCountsJob(List<CompletableFuture<Void>> work) {
         final KeyExpression recordCountKey = getRecordMetaData().getRecordCountKey();
-        if (recordCountKey == null) {
+        if (recordCountKey == null ||
+                getRecordStoreState().getStoreHeader().getRecordCountState() == RecordMetaDataProto.DataStoreInfo.RecordCountState.DISABLED) {
             return;
         }
         if (LOGGER.isDebugEnabled()) {

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

@@ -21,6 +21,7 @@
 package com.apple.foundationdb.record.provider.foundationdb;
 
 import com.apple.foundationdb.annotation.API;
+import com.apple.foundationdb.record.RecordMetaDataProto;
 
 import javax.annotation.Nonnull;
 import java.util.Arrays;
@@ -53,7 +54,7 @@ public enum FormatVersion implements Comparable<FormatVersion> {
      */
     INFO_ADDED(1),
     /**
-     * This FormatVersion introduces support for tracking record conuts as defined by:
+     * This FormatVersion introduces support for tracking record counts as defined by:
      * {@link com.apple.foundationdb.record.RecordMetaData#getRecordCountKey()}.
      */
     RECORD_COUNT_ADDED(2),
@@ -145,7 +146,12 @@ public enum FormatVersion implements Comparable<FormatVersion> {
     /**
      * This FormatVersion allows building non-idempotent indexes (e.g. COUNT) from a source index.
      */
-    CHECK_INDEX_BUILD_TYPE_DURING_UPDATE(10);
+    CHECK_INDEX_BUILD_TYPE_DURING_UPDATE(10),
+    /**
+     * This FormatVersion allows setting a state for the RecordCountKey on an individual store.
+     * @see RecordMetaDataProto.DataStoreInfo#getRecordCountState()
+     */
+    RECORD_COUNT_STATE(11);
 
     private final int value;
 

fdb-record-layer-core/src/main/proto/record_metadata.proto

@@ -95,6 +95,27 @@ message DataStoreInfo {
   // When considering using this, one should also consider introducing a new record type to store the data.
   // This was introduced with FormatVersion.HEADER_USER_FIELDS.
   repeated UserFieldEntry user_field = 8;
+
+  // The state of the RecordCount as defined by record_count_key. This is similar to IndexState, but since there is no
+  // mechanism for building the RecordCount across multiple transactions, the state transitions are slightly different.
+  enum RecordCountState {
+    // The RecordCountKey is maintained and usable for FDBRecordStore.getSnapshotRecordCount
+    READABLE = 1;
+    // The RecordCount as defined by record_count_key is maintained, but not queryable
+    // This can be freely updated back to READABLE
+    WRITE_ONLY = 2;
+    // The RecordCount as defined by record_count_key is not maintained.
+    // Once disabled, there is no way to go back to WriteOnly
+    DISABLED = 3;
+  }
+
+  // Whether the RecordCount as defined by record_count_key is usable by FDBRecordStore.getSnapshotRecordCount and
+  // whether it is maintained at all.
+  // This helps as a mechanism to de-risk removing the deprecated RecordCountKey from metadata by disabling it for
+  // queries first, and then disabling entirely, but only for specific stores. Once validated that the necessary COUNT
+  // indexes are working sufficiently, the RecordCountKey can be removed from the metadata more safely.
+  // This was introduced with FormatVersion.RECORD_COUNT_STATE
+  optional RecordCountState record_count_state = 9;
 }
 
 message Index {