addressing more comments

Author: normen662

fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java

@@ -21,6 +21,7 @@
 package com.apple.foundationdb.async.hnsw;
 
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -67,12 +68,12 @@ default void onNodeRead(int layer, @Nonnull Node<? extends NodeReference> node)
      * The default implementation is a no-op and does nothing.
      * @param layer the layer from which the key-value pair was read.
      * @param key the key that was read, guaranteed to be non-null.
-     * @param value the value associated with the key, guaranteed to be non-null.
+     * @param value the value associated with the key, can be null if the key was not found
      */
     @SuppressWarnings("unused")
     default void onKeyValueRead(int layer,
                                 @Nonnull byte[] key,
-                                @Nonnull byte[] value) {
+                                @Nullable byte[] value) {
         // nothing
     }
 }

fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java

@@ -235,7 +235,7 @@ static RealVector vectorFromBytes(@Nonnull final Config config, @Nonnull final b
      */
     @Nonnull
     @SuppressWarnings("PrimitiveArrayArgumentToVarargsMethod")
-    static Tuple tupleFromVector(final RealVector vector) {
+    static Tuple tupleFromVector(@Nonnull final RealVector vector) {
         return Tuple.from(vector.getRawData());
     }
 
@@ -249,15 +249,15 @@ static CompletableFuture<AccessInfo> fetchAccessInfo(@Nonnull final Config confi
                                                          @Nonnull final ReadTransaction readTransaction,
                                                          @Nonnull final Subspace subspace,
                                                          @Nonnull final OnReadListener onReadListener) {
-        final Subspace entryNodeSubspace = subspace.subspace(Tuple.from(SUBSPACE_PREFIX_ACCESS_INFO));
+        final Subspace entryNodeSubspace = accessInfoSubspace(subspace);
         final byte[] key = entryNodeSubspace.pack();
 
         return readTransaction.get(key)
                 .thenApply(valueBytes -> {
+                    onReadListener.onKeyValueRead(-1, key, valueBytes);
                     if (valueBytes == null) {
                         return null; // not a single node in the index
                     }
-                    onReadListener.onKeyValueRead(-1, key, valueBytes);
 
                     final Tuple entryTuple = Tuple.fromBytes(valueBytes);
                     final int layer = (int)entryTuple.getLong(0);
@@ -293,7 +293,7 @@ static void writeAccessInfo(@Nonnull final Transaction transaction,
                                 @Nonnull final Subspace subspace,
                                 @Nonnull final AccessInfo accessInfo,
                                 @Nonnull final OnWriteListener onWriteListener) {
-        final Subspace entryNodeSubspace = subspace.subspace(Tuple.from(SUBSPACE_PREFIX_ACCESS_INFO));
+        final Subspace entryNodeSubspace = accessInfoSubspace(subspace);
         final EntryNodeReference entryNodeReference = accessInfo.getEntryNodeReference();
         final RealVector centroid = accessInfo.getCentroid();
         final byte[] key = entryNodeSubspace.pack();
@@ -312,7 +312,7 @@ static CompletableFuture<List<AggregatedVector>> consumeSampledVectors(@Nonnull
                                                                            @Nonnull final Subspace subspace,
                                                                            final int numMaxVectors,
                                                                            @Nonnull final OnReadListener onReadListener) {
-        final Subspace prefixSubspace = subspace.subspace(Tuple.from(SUBSPACE_PREFIX_SAMPLES));
+        final Subspace prefixSubspace = samplesSubspace(subspace);
 
         final byte[] prefixKey = prefixSubspace.pack();
         final ReadTransaction snapshot = transaction.snapshot();
@@ -326,6 +326,7 @@ static CompletableFuture<List<AggregatedVector>> consumeSampledVectors(@Nonnull
                         final byte[] key = keyValue.getKey();
                         final byte[] value = keyValue.getValue();
                         resultBuilder.add(aggregatedVectorFromRaw(prefixSubspace, key, value));
+                        transaction.addReadConflictKey(key);
                         transaction.clear(key);
                         onReadListener.onKeyValueRead(-1, key, value);
                     }
@@ -338,7 +339,7 @@ static void appendSampledVector(@Nonnull final Transaction transaction,
                                     final int partialCount,
                                     @Nonnull final Transformed<RealVector> vector,
                                     @Nonnull final OnWriteListener onWriteListener) {
-        final Subspace prefixSubspace = subspace.subspace(Tuple.from(SUBSPACE_PREFIX_SAMPLES));
+        final Subspace prefixSubspace = samplesSubspace(subspace);
         final Subspace keySubspace = prefixSubspace.subspace(Tuple.from(partialCount, UUID.randomUUID()));
         final byte[] prefixKey = keySubspace.pack();
         // getting underlying is okay as it is only written to the database
@@ -348,8 +349,7 @@ static void appendSampledVector(@Nonnull final Transaction transaction,
     }
 
     static void removeAllSampledVectors(@Nonnull final Transaction transaction, @Nonnull final Subspace subspace) {
-        final Subspace prefixSubspace =
-                subspace.subspace(Tuple.from(SUBSPACE_PREFIX_SAMPLES));
+        final Subspace prefixSubspace = samplesSubspace(subspace);
 
         final byte[] prefixKey = prefixSubspace.pack();
         final Range range = Range.startsWith(prefixKey);
@@ -366,4 +366,14 @@ private static AggregatedVector aggregatedVectorFromRaw(@Nonnull final Subspace
 
         return new AggregatedVector(partialCount, AffineOperator.identity().transform(vector));
     }
+
+    @Nonnull
+    static Subspace accessInfoSubspace(@Nonnull final Subspace rootSubspace) {
+        return rootSubspace.subspace(Tuple.from(SUBSPACE_PREFIX_ACCESS_INFO));
+    }
+
+    @Nonnull
+    static Subspace samplesSubspace(@Nonnull final Subspace rootSubspace) {
+        return rootSubspace.subspace(Tuple.from(SUBSPACE_PREFIX_SAMPLES));
+    }
 }

fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWTest.java

@@ -59,6 +59,7 @@
 import org.slf4j.LoggerFactory;
 
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import java.io.IOException;
 import java.nio.channels.FileChannel;
 import java.nio.file.Path;
@@ -495,9 +496,9 @@ public void onNodeRead(final int layer, @Nonnull final Node<? extends NodeRefere
         }
 
         @Override
-        public void onKeyValueRead(final int layer, @Nonnull final byte[] key, @Nonnull final byte[] value) {
+        public void onKeyValueRead(final int layer, @Nonnull final byte[] key, @Nullable final byte[] value) {
             bytesReadByLayer.compute(layer, (l, oldValue) -> (oldValue == null ? 0 : oldValue) +
-                    key.length + value.length);
+                    key.length + (value == null ? 0 : value.length));
         }
     }