Fix bug of scan aggregate index returning empty non-end continuation

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

@@ -41,9 +41,10 @@ private IndexPrefetchRangeKeyValueCursor(@Nonnull final FDBRecordContext context
                                              @Nonnull final AsyncIterator<MappedKeyValue> iterator,
                                              int prefixLength,
                                              @Nonnull final CursorLimitManager limitManager,
-                                             int valuesLimit) {
+                                             int valuesLimit,
+                                             @Nonnull SerializationMode serializationMode) {
 
-        super(context, iterator, prefixLength, limitManager, valuesLimit);
+        super(context, iterator, prefixLength, limitManager, valuesLimit, serializationMode);
     }
 
     /**
@@ -69,7 +70,7 @@ public IndexPrefetchRangeKeyValueCursor build() {
             AsyncIterator<MappedKeyValue> iterator = getTransaction()
                     .getMappedRange(getBegin(), getEnd(), mapper, getLimit(), isReverse(), getStreamingMode())
                     .iterator();
-            return new IndexPrefetchRangeKeyValueCursor(getContext(), iterator, getPrefixLength(), getLimitManager(), getValuesLimit());
+            return new IndexPrefetchRangeKeyValueCursor(getContext(), iterator, getPrefixLength(), getLimitManager(), getValuesLimit(), serializationMode);
         }
 
         @Override

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

@@ -38,8 +38,9 @@ private KeyValueCursor(@Nonnull final FDBRecordContext context,
                              @Nonnull final AsyncIterator<KeyValue> iterator,
                              int prefixLength,
                              @Nonnull final CursorLimitManager limitManager,
-                             int valuesLimit) {
-        super(context, iterator, prefixLength, limitManager, valuesLimit);
+                             int valuesLimit,
+                             @Nonnull SerializationMode serializationMode) {
+        super(context, iterator, prefixLength, limitManager, valuesLimit, serializationMode);
     }
 
     /**
@@ -77,7 +78,7 @@ public KeyValueCursor build() {
             final AsyncIterator<KeyValue> iterator = getTransaction()
                     .getRange(getBegin(), getEnd(), getLimit(), isReverse(), getStreamingMode())
                     .iterator();
-            return new KeyValueCursor(getContext(), iterator, getPrefixLength(), getLimitManager(), getValuesLimit());
+            return new KeyValueCursor(getContext(), iterator, getPrefixLength(), getLimitManager(), getValuesLimit(), serializationMode);
         }
 
         @Override

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

@@ -33,6 +33,7 @@
 import com.apple.foundationdb.record.KeyRange;
 import com.apple.foundationdb.record.RecordCoreException;
 import com.apple.foundationdb.record.RecordCursorContinuation;
+import com.apple.foundationdb.record.RecordCursorProto;
 import com.apple.foundationdb.record.RecordCursorResult;
 import com.apple.foundationdb.record.ScanProperties;
 import com.apple.foundationdb.record.TupleRange;
@@ -42,11 +43,13 @@
 import com.apple.foundationdb.subspace.Subspace;
 import com.apple.foundationdb.tuple.Tuple;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.ZeroCopyByteString;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.util.Arrays;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
@@ -61,22 +64,26 @@ public abstract class KeyValueCursorBase<K extends KeyValue> extends AsyncIterat
     private final int prefixLength;
     @Nonnull
     private final CursorLimitManager limitManager;
-    private int valuesLimit;
+    private final int valuesLimit;
     // the pointer may be mutated, but the actual array must never be mutated or continuations will break
     @Nullable
     private byte[] lastKey;
+    @Nonnull
+    private final SerializationMode serializationMode;
 
     protected KeyValueCursorBase(@Nonnull final FDBRecordContext context,
                                  @Nonnull final AsyncIterator<K> iterator,
                                  int prefixLength,
                                  @Nonnull final CursorLimitManager limitManager,
-                                 int valuesLimit) {
+                                 int valuesLimit,
+                                 @Nonnull final SerializationMode serializationMode) {
         super(context.getExecutor(), iterator);
 
         this.context = context;
         this.prefixLength = prefixLength;
         this.limitManager = limitManager;
         this.valuesLimit = valuesLimit;
+        this.serializationMode = serializationMode;
 
         context.instrument(FDBStoreTimer.DetailEvents.GET_SCAN_RANGE_RAW_FIRST_CHUNK, iterator.onHasNext());
     }
@@ -131,21 +138,23 @@ public RecordCursorResult<K> getNext() {
 
     @Nonnull
     private RecordCursorContinuation continuationHelper() {
-        return new Continuation(lastKey, prefixLength);
+        return new Continuation(lastKey, prefixLength, serializationMode);
     }
 
-    private static class Continuation implements RecordCursorContinuation {
+    public static class Continuation implements RecordCursorContinuation {
         @Nullable
         private final byte[] lastKey;
         private final int prefixLength;
+        private final SerializationMode serializationMode;
 
-        public Continuation(@Nullable final byte[] lastKey, final int prefixLength) {
+        public Continuation(@Nullable final byte[] lastKey, final int prefixLength, final SerializationMode serializationMode) {
             // Note that doing this without a full copy is dangerous if the array is ever mutated.
             // Currently, this never happens and the only thing that changes is which array lastKey points to.
             // However, if logic in KeyValueCursor or KeyValue changes, this could break continuations.
             // To resolve it, we could resort to doing a full copy here, although that's somewhat expensive.
             this.lastKey = lastKey;
             this.prefixLength = prefixLength;
+            this.serializationMode = serializationMode;
         }
 
         @Override
@@ -156,21 +165,83 @@ public boolean isEnd() {
         @Nonnull
         @Override
         public ByteString toByteString() {
-            if (lastKey == null) {
-                return ByteString.EMPTY;
+            if (serializationMode == SerializationMode.TO_OLD) {
+                // lastKey = null when source iterator hit limit that we passed down.
+                if (lastKey == null) {
+                    return ByteString.EMPTY;
+                }
+                ByteString base = ZeroCopyByteString.wrap(lastKey);
+                // when prefixLength == lastKey.length, toByteString() also returns ByteString.EMPTY
+                return base.substring(prefixLength, lastKey.length);
+            } else {
+                return toProto().toByteString();
             }
-            ByteString base = ZeroCopyByteString.wrap(lastKey);
-            return base.substring(prefixLength, lastKey.length);
         }
 
         @Nullable
         @Override
         public byte[] toBytes() {
+            if (lastKey == null) {
+                return null;
+            }
+            ByteString byteString = toByteString();
+            return byteString.isEmpty() ? new byte[0] : byteString.toByteArray();
+        }
+
+        @Nullable
+        public byte[] getInnerContinuationInBytes() {
             if (lastKey == null) {
                 return null;
             }
             return Arrays.copyOfRange(lastKey, prefixLength, lastKey.length);
         }
+
+        @Nonnull
+        public ByteString getInnerContinuationInByteString() {
+            if (lastKey == null) {
+                return ByteString.EMPTY;
+            }
+            ByteString base = ZeroCopyByteString.wrap(lastKey);
+            return base.substring(prefixLength, lastKey.length);
+        }
+
+        public static byte[] fromRawBytes(@Nullable byte[] rawBytes, SerializationMode serializationMode) {
+            if (rawBytes == null) {
+                return null;
+            }
+            if (serializationMode == SerializationMode.TO_OLD) {
+                return rawBytes;
+            }
+            try {
+                RecordCursorProto.KeyValueCursorContinuation continuationProto = RecordCursorProto.KeyValueCursorContinuation.parseFrom(rawBytes);
+                if (continuationProto.hasPrefixLength()) {
+                    return continuationProto.getContinuation().toByteArray();
+                } else {
+                    // parseFrom can parse an old serialization result as the new proto, wrong deserialization
+                    return rawBytes;
+                }
+            } catch (InvalidProtocolBufferException ipbe) {
+                return rawBytes;
+            }
+        }
+
+        @Nonnull
+        private RecordCursorProto.KeyValueCursorContinuation toProto() {
+            RecordCursorProto.KeyValueCursorContinuation.Builder builder = RecordCursorProto.KeyValueCursorContinuation.newBuilder();
+            if (lastKey == null) {
+                // proto.hasContinuation() = false when lastKey = null
+                return builder.setPrefixLength(prefixLength).build();
+            } else {
+                ByteString base = ZeroCopyByteString.wrap(Objects.requireNonNull(lastKey));
+                // proto.hasContinuation() = ByteString.EMPTY when prefixLength = lastKey.length
+                return builder.setContinuation(base.substring(prefixLength, lastKey.length)).setPrefixLength(prefixLength).build();
+            }
+        }
+    }
+
+    public enum SerializationMode {
+        TO_OLD,
+        TO_NEW
     }
 
     /**
@@ -208,9 +279,11 @@ public abstract static class Builder<T extends Builder<T>> {
         private StreamingMode streamingMode;
         private KeySelector begin;
         private KeySelector end;
+        protected SerializationMode serializationMode;
 
         protected Builder(@Nonnull Subspace subspace) {
             this.subspace = subspace;
+            this.serializationMode = SerializationMode.TO_OLD;
         }
 
         /**
@@ -247,10 +320,12 @@ protected void prepare() {
             prefixLength = calculatePrefixLength();
 
             reverse = scanProperties.isReverse();
+
             if (continuation != null) {
-                final byte[] continuationBytes = new byte[prefixLength + continuation.length];
+                byte[] realContinuation = KeyValueCursorBase.Continuation.fromRawBytes(continuation, serializationMode);
+                final byte[] continuationBytes = new byte[prefixLength + realContinuation.length];
                 System.arraycopy(lowBytes, 0, continuationBytes, 0, prefixLength);
-                System.arraycopy(continuation, 0, continuationBytes, prefixLength, continuation.length);
+                System.arraycopy(realContinuation, 0, continuationBytes, prefixLength, realContinuation.length);
                 if (reverse) {
                     highBytes = continuationBytes;
                     highEndpoint = EndpointType.CONTINUATION;
@@ -334,6 +409,11 @@ public T setHigh(@Nonnull byte[] highBytes, @Nonnull EndpointType highEndpoint)
             return self();
         }
 
+        public T setSerializationMode(@Nonnull final SerializationMode serializationMode) {
+            this.serializationMode = serializationMode;
+            return self();
+        }
+
         /**
          * Calculate the key prefix length for the returned values. This will be used to derive the primary key used in
          * the calculated continuation.