NoSQL: simplify node allocation

Author: snazy

persistence/nosql/nodes/impl/src/main/java/org/apache/polaris/nodes/impl/NodeManagementImpl.java

@@ -23,6 +23,7 @@
 import static java.lang.Integer.bitCount;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.primitives.Ints;
 import jakarta.annotation.Nonnull;
 import jakarta.annotation.PreDestroy;
 import jakarta.enterprise.context.ApplicationScoped;
@@ -35,15 +36,14 @@
 import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.BitSet;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
-import java.util.function.IntSupplier;
+import java.util.stream.IntStream;
 import org.apache.polaris.ids.api.IdGenerator;
 import org.apache.polaris.ids.api.IdGeneratorSpec;
 import org.apache.polaris.ids.api.ImmutableIdGeneratorSpec;
@@ -255,35 +255,10 @@ private LeaseParams leaseInternal() {
 
     var now = clock.currentInstant();
 
-    var checkedIds = new BitSet(numNodeIds); // 128 bytes for 1024 nodes - not too much
-    var rand = ThreadLocalRandom.current();
-
-    var generateNodeId =
-        (IntSupplier)
-            () -> {
-              if (checkedIds.cardinality() == numNodeIds) {
-                return -1;
-              }
-              while (true) {
-                // The implementation _could_ use BitSet.nextClearBit(), but this is intentionally
-                // _not_ done here to let two concurrently started nodes not generate the same set
-                // of "fallback" node IDs to use (after checking with a hash over the network
-                // interface).
-                var nodeId = rand.nextInt(numNodeIds);
-                if (!checkedIds.get(nodeId)) {
-                  checkedIds.set(nodeId);
-                  return nodeId;
-                }
-              }
-            };
-
-    var ids = new Integer[CHECK_BATCH_SIZE];
-
-    // First, try with a hash of the local network address...
+    // First, try with a hash of the local network address.
     // The node ID in this attempt is somewhat deterministic for the local machine.
-    // This is not really a necessity, but it can reduce the overall number of ever allocated node
-    // IDs. It is rather a cosmetic thing, so that a LIST NODES command in the Persistence-REPL
-    // doesn't flood users after a couple of restarts.
+    // This is not really necessary, but it can reduce the overall number of ever-allocated node
+    // IDs.
     try {
       var hashOverNetworkInterfaces =
           NetworkInterface.networkInterfaces()
@@ -309,39 +284,22 @@ private LeaseParams leaseInternal() {
 
       var nodeId = hashOverNetworkInterfaces & (numNodeIds - 1);
 
-      var leased = tryLeaseFromCandidates(new Integer[] {nodeId}, now);
+      var leased = tryLeaseFromCandidates(new int[] {nodeId}, now);
       if (leased != null) {
         return leased;
       }
     } catch (SocketException e) {
       // ignore
     }
 
-    // Next, try with randomly picked node IDs and attempt to lease...
-    for (int i = 0; i < 3 * numNodeIds / CHECK_BATCH_SIZE; i++) {
-      Arrays.fill(ids, null);
-      for (int j = 0; j < ids.length; j++) {
-        var nodeId = generateNodeId.getAsInt();
-        if (nodeId == -1) {
-          break;
-        }
-        ids[j] = nodeId;
-      }
+    // If the lease-attempt using the hash over network-interfaced did not succeed, try with
+    // randomly picked node IDs.
 
-      var leased = tryLeaseFromCandidates(ids, now);
-      if (leased != null) {
-        return leased;
-      }
-    }
+    var nodeIdsToCheck = IntStream.range(0, numNodeIds).toArray();
+    Ints.rotate(nodeIdsToCheck, ThreadLocalRandom.current().nextInt(numNodeIds));
 
-    // Next, try again, but using sequential nodeIds starting at a random offset...
-    var offset = rand.nextInt(numNodeIds);
     for (int i = 0; i < numNodeIds; i += CHECK_BATCH_SIZE) {
-      Arrays.fill(ids, null);
-      for (int j = 0; j < ids.length; j++) {
-        var nodeId = (offset + i + j) % numNodeIds;
-        ids[j] = nodeId;
-      }
+      var ids = Arrays.copyOfRange(nodeIdsToCheck, i, i + CHECK_BATCH_SIZE);
       var leased = tryLeaseFromCandidates(ids, now);
       if (leased != null) {
         return leased;
@@ -376,19 +334,16 @@ public void close() {
     }
   }
 
-  private LeaseParams tryLeaseFromCandidates(Integer[] nodeIds, Instant now) {
+  private LeaseParams tryLeaseFromCandidates(int[] nodeIds, Instant now) {
     var nodesFetched = nodeStore.fetchMany(nodeIds);
     for (int i = 0; i < nodeIds.length; i++) {
       // NodeStore.fetchMany MUST return as many elements as the input array.
       var nodeId = nodeIds[i];
-      if (nodeId != null) {
-        var id = (int) nodeId;
-        var node = nodesFetched[i];
-        if (canLease(node, now)) {
-          var leased = tryLease(id, node);
-          if (leased != null) {
-            return leased;
-          }
+      var node = nodesFetched[i];
+      if (canLease(node, now)) {
+        var leased = tryLease(nodeId, node);
+        if (leased != null) {
+          return leased;
         }
       }
     }

persistence/nosql/nodes/impl/src/test/java/org/apache/polaris/nodes/impl/TestNodeManagementImpl.java

@@ -154,6 +154,7 @@ public void simple() {
                     .build())
             .build();
     try (var mgmt = new NodeManagementImpl(config, clock, new MockNodeStoreFactory(), scheduler)) {
+      soft.assertThat(mgmt.maxNumberOfNodes()).isEqualTo(config.numNodes());
       var lease = mgmt.lease();
       soft.assertThat(lease).isNotNull();
       soft.assertThat(lease.nodeIdIfValid()).isNotEqualTo(-1);

persistence/nosql/nodes/impl/src/testFixtures/java/org/apache/polaris/nodes/impl/MockNodeStore.java

@@ -33,7 +33,7 @@ public class MockNodeStore implements NodeStore {
   @Override
   public NodeState persist(
       int nodeId, Optional<NodeState> expectedNodeState, @Nonnull NodeState newState) {
-    if (expectedNodeState.isPresent()) {
+    if (nodeId >= 0 && expectedNodeState.isPresent()) {
       var result =
           nodeStates.computeIfPresent(
               nodeId,
@@ -56,11 +56,11 @@ public NodeState persist(
 
   @Nonnull
   @Override
-  public NodeState[] fetchMany(@Nonnull Integer... nodeIds) {
+  public NodeState[] fetchMany(@Nonnull int... nodeIds) {
     var r = new NodeState[nodeIds.length];
     for (int i = 0; i < nodeIds.length; i++) {
       var id = nodeIds[i];
-      if (id != null) {
+      if (id >= 0) {
         r[i] = nodeStates.get(id);
       }
     }

persistence/nosql/nodes/spi/src/main/java/org/apache/polaris/nodes/spi/NodeStore.java

@@ -27,7 +27,7 @@ public interface NodeStore {
   Optional<NodeState> fetch(int nodeId);
 
   @Nonnull
-  NodeState[] fetchMany(@Nonnull Integer... nodeIds);
+  NodeState[] fetchMany(@Nonnull int... nodeIds);
 
   @Nullable
   NodeState persist(int nodeId, Optional<NodeState> expectedNodeState, @Nonnull NodeState newState);