Fix: Allow reserving threads in Java

ashvardanian · ashvardanian · commit 55b5f7655227 · 2025-09-03T14:11:29.000Z
diff --git a/java/cloud/unum/usearch/Index.java b/java/cloud/unum/usearch/Index.java
@@ -274,7 +274,34 @@ public void reserve(long capacity) {
         if (c_ptr == 0) {
             throw new IllegalStateException("Index already closed");
         }
-        c_reserve(c_ptr, capacity);
+        // Pass zeros to use all available contexts on the device
+        c_reserve(c_ptr, capacity, 0, 0);
+    }
+
+    /**
+     * Reserves memory and configures thread contexts.
+     * Use this to explicitly control concurrent add/search capacities.
+     *
+     * @param capacity desired total capacity
+     * @param threadsAdd maximum concurrent add contexts
+     * @param threadsSearch maximum concurrent search contexts
+     */
+    public void reserve(long capacity, long threadsAdd, long threadsSearch) {
+        if (c_ptr == 0) {
+            throw new IllegalStateException("Index already closed");
+        }
+        c_reserve(c_ptr, capacity, threadsAdd, threadsSearch);
+    }
+
+    /**
+     * Reserves memory and sets the same number of contexts
+     * for both add and search operations.
+     *
+     * @param capacity desired total capacity
+     * @param threads number of contexts for both add and search
+     */
+    public void reserve(long capacity, long threads) {
+        reserve(capacity, threads, threads);
     }
 
     /**
@@ -1021,7 +1048,7 @@ private static native long c_create(
 
     private static native long c_capacity(long ptr);
 
-    private static native void c_reserve(long ptr, long capacity);
+    private static native void c_reserve(long ptr, long capacity, long threadsAdd, long threadsSearch);
 
     private static native void c_save(long ptr, String path);
 
diff --git a/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp b/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp
@@ -143,8 +143,22 @@ JNIEXPORT jlong JNICALL Java_cloud_unum_usearch_Index_c_1capacity(JNIEnv*, jclas
     return reinterpret_cast<index_dense_t*>(c_ptr)->capacity();
 }
 
-JNIEXPORT void JNICALL Java_cloud_unum_usearch_Index_c_1reserve(JNIEnv* env, jclass, jlong c_ptr, jlong capacity) {
-    if (!reinterpret_cast<index_dense_t*>(c_ptr)->try_reserve(static_cast<std::size_t>(capacity))) {
+JNIEXPORT void JNICALL Java_cloud_unum_usearch_Index_c_1reserve(JNIEnv* env, jclass, jlong c_ptr, jlong capacity,
+                                                                jlong threads_add, jlong threads_search) {
+    std::size_t t_add = static_cast<std::size_t>(threads_add);
+    std::size_t t_search = static_cast<std::size_t>(threads_search);
+    if (t_add == 0 || t_search == 0) {
+        std::size_t hc = std::thread::hardware_concurrency();
+        if (hc == 0)
+            hc = 1; // fallback to 1 if the runtime can't  report
+        if (t_add == 0)
+            t_add = hc;
+        if (t_search == 0)
+            t_search = hc;
+    }
+    index_limits_t limits(static_cast<std::size_t>(capacity), t_add);
+    limits.threads_search = t_search;
+    if (!reinterpret_cast<index_dense_t*>(c_ptr)->try_reserve(limits)) {
         jclass jc = (*env).FindClass("java/lang/Error");
         if (jc)
             (*env).ThrowNew(jc, "Failed to grow vector index!");
diff --git a/java/cloud/unum/usearch/cloud_unum_usearch_Index.h b/java/cloud/unum/usearch/cloud_unum_usearch_Index.h
diff --git a/java/test/IndexTest.java b/java/test/IndexTest.java
@@ -267,13 +267,14 @@ public void testGetIntoBufferMethods() {
     @Test
     public void testConcurrentAdd() throws Exception {
         try (Index index = new Index.Config().metric("cos").dimensions(4).build()) {
-            index.reserve(1000);
+            final int threadsCount = 10;
+            index.reserve(1000, threadsCount);
 
-            ExecutorService executor = Executors.newFixedThreadPool(10);
+            ExecutorService executor = Executors.newFixedThreadPool(threadsCount);
             @SuppressWarnings("unchecked")
-            CompletableFuture<Void>[] futures = new CompletableFuture[10];
+            CompletableFuture<Void>[] futures = new CompletableFuture[threadsCount];
 
-            for (int t = 0; t < 10; t++) {
+            for (int t = 0; t < threadsCount; t++) {
                 final int threadId = t;
                 futures[t]
                         = CompletableFuture.runAsync(
@@ -290,25 +291,26 @@ public void testConcurrentAdd() throws Exception {
             CompletableFuture.allOf(futures).get(10, TimeUnit.SECONDS);
             executor.shutdown();
 
-            assertEquals(500, index.size());
+            assertEquals(50L * threadsCount, index.size());
         }
     }
 
     @Test
     public void testConcurrentSearch() throws Exception {
         try (Index index = new Index.Config().metric("cos").dimensions(4).build()) {
-            index.reserve(100);
+            final int threadsCount = 5;
+            index.reserve(100, threadsCount);
 
             // Add some vectors first
             for (int i = 0; i < 100; i++) {
                 index.add(i, randomVector(4));
             }
 
-            ExecutorService executor = Executors.newFixedThreadPool(5);
+            ExecutorService executor = Executors.newFixedThreadPool(threadsCount);
             @SuppressWarnings("unchecked")
-            CompletableFuture<long[]>[] futures = new CompletableFuture[5];
+            CompletableFuture<long[]>[] futures = new CompletableFuture[threadsCount];
 
-            for (int t = 0; t < 5; t++) {
+            for (int t = 0; t < threadsCount; t++) {
                 futures[t]
                         = CompletableFuture.supplyAsync(
                                 () -> {
@@ -328,56 +330,6 @@ public void testConcurrentSearch() throws Exception {
         }
     }
 
-    @Test
-    public void testMixedConcurrency() throws Exception {
-        try (Index index = new Index.Config().metric("cos").dimensions(3).build()) {
-            index.reserve(200);
-
-            ExecutorService executor = Executors.newFixedThreadPool(8);
-            @SuppressWarnings("unchecked")
-            CompletableFuture<Void>[] addFutures = new CompletableFuture[4];
-            @SuppressWarnings("unchecked")
-            CompletableFuture<Void>[] searchFutures = new CompletableFuture[4];
-
-            // Add operations
-            for (int t = 0; t < 4; t++) {
-                final int threadId = t;
-                addFutures[t]
-                        = CompletableFuture.runAsync(
-                                () -> {
-                                    for (int i = 0; i < 30; i++) {
-                                        long key = threadId * 30L + i;
-                                        index.add(key, randomVector(3));
-                                    }
-                                },
-                                executor);
-            }
-
-            // Wait for some adds to complete, then start searches
-            Thread.sleep(100);
-
-            // Search operations
-            for (int t = 0; t < 4; t++) {
-                searchFutures[t]
-                        = CompletableFuture.runAsync(
-                                () -> {
-                                    for (int i = 0; i < 10; i++) {
-                                        float[] queryVector = randomVector(3);
-                                        long[] results = index.search(queryVector, 5);
-                                        assertTrue(results.length >= 0);
-                                    }
-                                },
-                                executor);
-            }
-
-            CompletableFuture.allOf(addFutures).get(15, TimeUnit.SECONDS);
-            CompletableFuture.allOf(searchFutures).get(15, TimeUnit.SECONDS);
-            executor.shutdown();
-
-            assertEquals(120, index.size());
-        }
-    }
-
     @Test
     public void testBatchAdd() {
         try (Index index = new Index.Config().metric("cos").dimensions(2).build()) {
@@ -751,30 +703,30 @@ public void testPlatformCapabilities() {
         String[] available = Index.hardwareAccelerationAvailable();
         assertNotEquals("Available capabilities should not be null", null, available);
         assertTrue("Platform should have at least serial capability", available.length > 0);
-        
+
         // Test compile-time capabilities
         String[] compiled = Index.hardwareAccelerationCompiled();
         assertNotEquals("Compiled capabilities should not be null", null, compiled);
         assertTrue("Should have at least serial compiled", compiled.length > 0);
-        
+
         // Should always include serial as baseline in both
         boolean hasAvailableSerial = false;
         boolean hasCompiledSerial = false;
-        
+
         for (String cap : available) {
             if ("serial".equals(cap)) {
                 hasAvailableSerial = true;
                 break;
             }
         }
-        
+
         for (String cap : compiled) {
             if ("serial".equals(cap)) {
                 hasCompiledSerial = true;
                 break;
             }
         }
-        
+
         assertTrue("Platform should always support serial capability", hasAvailableSerial);
         assertTrue("Serial should always be compiled", hasCompiledSerial);
 
