various fixes

Author: normen662

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

@@ -86,7 +86,8 @@ public CompactStorageAdapter(@Nonnull final Config config,
      * @param layer the layer of the node to fetch
      * @param primaryKey the primary key of the node to fetch
      *
-     * @return a future that will complete with the fetched {@link AbstractNode}
+     * @return a future that will complete with the fetched {@link AbstractNode} or {@code null} if the node cannot
+     *         be fetched
      *
      * @throws IllegalStateException if the node cannot be found in the database for the given key
      */
@@ -101,7 +102,7 @@ protected CompletableFuture<AbstractNode<NodeReference>> fetchNodeInternal(@Nonn
         return readTransaction.get(keyBytes)
                 .thenApply(valueBytes -> {
                     if (valueBytes == null) {
-                        throw new IllegalStateException("cannot fetch node");
+                        return null;
                     }
                     return nodeFromRaw(storageTransform, layer, primaryKey, keyBytes, valueBytes);
                 });

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

@@ -581,7 +581,7 @@ private Quantizer quantizer(@Nullable final AccessInfo accessInfo) {
         return onReadListener.onAsyncRead(
                         storageAdapter.fetchNode(readTransaction, storageTransform, layer,
                                 nodeReference.getPrimaryKey()))
-                .thenApply(node -> biMapFunction.apply(nodeReference, node));
+                .thenApply(node -> biMapFunction.apply(nodeReference, Objects.requireNonNull(node)));
     }
 
     /**
@@ -754,13 +754,28 @@ public CompletableFuture<Void> insert(@Nonnull final Transaction transaction, @N
         }
 
         return StorageAdapter.fetchAccessInfo(getConfig(), transaction, getSubspace(), getOnReadListener())
-                .thenCompose(accessInfo -> {
-                    final AccessInfo currentAccessInfo;
+                .thenCombine(exists(transaction, newPrimaryKey),
+                        (accessInfo, nodeAlreadyExists) -> {
+                            if (nodeAlreadyExists) {
+                                if (logger.isDebugEnabled()) {
+                                    logger.debug("new record already exists in HNSW with key={} on layer={}",
+                                            newPrimaryKey, insertionLayer);
+                                }
+                            }
+                            return new AccessInfoAndNodeExistence(accessInfo, nodeAlreadyExists);
+                        })
+                .thenCompose(accessInfoAndNodeExistence -> {
+                    if (accessInfoAndNodeExistence.isNodeExists()) {
+                        return AsyncUtil.DONE;
+                    }
+
+                    final AccessInfo accessInfo = accessInfoAndNodeExistence.getAccessInfo();
                     final AffineOperator storageTransform = storageTransform(accessInfo);
                     final Transformed<RealVector> transformedNewVector = storageTransform.transform(newVector);
                     final Quantizer quantizer = quantizer(accessInfo);
                     final Estimator estimator = quantizer.estimator();
 
+                    final AccessInfo currentAccessInfo;
                     if (accessInfo == null) {
                         // this is the first node
                         writeLonelyNodes(quantizer, transaction, newPrimaryKey, transformedNewVector,
@@ -821,6 +836,15 @@ public CompletableFuture<Void> insert(@Nonnull final Transaction transaction, @N
                 }).thenCompose(ignored -> AsyncUtil.DONE);
     }
 
+    @Nonnull
+    @VisibleForTesting
+    CompletableFuture<Boolean> exists(@Nonnull final ReadTransaction readTransaction,
+                                      @Nonnull final Tuple primaryKey) {
+        final StorageAdapter<? extends NodeReference> storageAdapter = getStorageAdapterForLayer(0);
+        return storageAdapter.fetchNode(readTransaction, AffineOperator.identity(), 0, primaryKey)
+                .thenApply(Objects::nonNull);
+    }
+
     /**
      * Method to keep stats if necessary. Stats need to be kept and maintained when the client would like to use
      * e.g. RaBitQ as RaBitQ needs a stable somewhat correct centroid in order to function properly.
@@ -1546,4 +1570,24 @@ private static <T> List<T> drain(@Nonnull Queue<T> queue) {
         }
         return resultBuilder.build();
     }
+
+    private static class AccessInfoAndNodeExistence {
+        @Nullable
+        private final AccessInfo accessInfo;
+        private final boolean nodeExists;
+
+        public AccessInfoAndNodeExistence(@Nullable final AccessInfo accessInfo, final boolean nodeExists) {
+            this.accessInfo = accessInfo;
+            this.nodeExists = nodeExists;
+        }
+
+        @Nullable
+        public AccessInfo getAccessInfo() {
+            return accessInfo;
+        }
+
+        public boolean isNodeExists() {
+            return nodeExists;
+        }
+    }
 }