withMemoryLimitMB -> withArrowMemoryLimitMB

Author: sgrebnov

README.md

@@ -128,7 +128,7 @@ The `SpiceClient` uses an Arrow `RootAllocator` for managing off-heap memory. By
 
 ```java
 SpiceClient client = SpiceClient.builder()
-    .withMemoryLimitMB(1024) // 1GB limit
+    .withArrowMemoryLimitMB(1024) // 1GB limit
     .build();
 ```
 

src/main/java/ai/spice/SpiceClientBuilder.java

@@ -143,19 +143,26 @@ public SpiceClientBuilder withMaxRetries(int maxRetries) {
     /**
      * Sets the memory limit for Apache Arrow allocator in megabytes.
      * This controls the maximum amount of off-heap memory that can be allocated
-     * for Arrow Flight operations. If not set, the allocator will use all available memory.
+     * for Arrow Flight operations. If not set, the allocator will use all available
+     * memory.
      *
-     * @param memoryLimitMB Maximum memory limit in megabytes. Default is all available memory.
+     * @param memoryLimitMB Maximum memory limit in megabytes. Default is all
+     *                      available memory.
      *                      Must be positive.
      * @return The current instance of SpiceClientBuilder for method chaining.
      * @throws IllegalArgumentException if memoryLimitMB is not positive
      *
      * @see org.apache.arrow.memory.RootAllocator
      */
-    public SpiceClientBuilder withMemoryLimitMB(long memoryLimitMB) {
+    public SpiceClientBuilder withArrowMemoryLimitMB(long memoryLimitMB) {
         if (memoryLimitMB <= 0) {
             throw new IllegalArgumentException("Memory limit must be positive, got: " + memoryLimitMB + " MB");
         }
+        if (memoryLimitMB > Long.MAX_VALUE / 1024L / 1024L) {
+            throw new IllegalArgumentException(
+                    "Memory limit is too large: " + memoryLimitMB + " MB");
+        }
+
         this.memoryLimitMB = memoryLimitMB;
         return this;
     }

src/test/java/ai/spice/MemoryConfigurationTest.java

@@ -11,41 +11,70 @@ public void testMemoryLimitMBConfiguration() throws Exception {
         try {
             // Test that memory limit in MB is properly configured
             SpiceClient client = SpiceClient.builder()
-                .withMemoryLimitMB(128) // 128 MB
-                .build();
-            
+                    .withArrowMemoryLimitMB(128) // 128 MB
+                    .build();
+
             // If we reach here without exception, the configuration worked
             assertTrue("Memory configuration should not throw exception", true);
-            
+
             client.close();
         } catch (Exception e) {
-            // We expect some exceptions due to no local Spice instance,
-            // but not IllegalArgumentException from memory configuration
-            assertFalse("Should not throw IllegalArgumentException for valid memory config", 
-                       e instanceof IllegalArgumentException);
+            fail("Should not throw exception: " + e.getMessage());
         }
     }
 
     public void testInvalidMemoryLimitMB() throws Exception {
         try {
             SpiceClient.builder()
-                .withMemoryLimitMB(0) // Invalid: must be positive
-                .build();
+                    .withArrowMemoryLimitMB(0) // Invalid: must be positive
+                    .build();
             fail("Should throw IllegalArgumentException for zero memory limit");
         } catch (IllegalArgumentException e) {
             // Expected exception
-           assertTrue("Should throw IllegalArgumentException for zero memory limit", true);
+            assertTrue("Should throw IllegalArgumentException for zero memory limit", true);
         }
     }
 
     public void testNegativeMemoryLimitMB() throws Exception {
         try {
             SpiceClient.builder()
-                .withMemoryLimitMB(-100) // Invalid: negative
-                .build();
+                    .withArrowMemoryLimitMB(-100) // Invalid: negative
+                    .build();
             fail("Should throw IllegalArgumentException for negative memory limit");
         } catch (IllegalArgumentException e) {
             assertTrue("Should throw IllegalArgumentException for negative memory limit", true);
         }
     }
+
+    public void testOverflowProtection() throws Exception {
+        // Test overflow protection - values that would cause overflow when converted to
+        // bytes
+        long maxSafeMB = Long.MAX_VALUE / (1024L * 1024L);
+        long overflowValue = maxSafeMB + 1;
+
+        try {
+            SpiceClient.builder()
+                    .withArrowMemoryLimitMB(overflowValue) // Would cause overflow
+                    .build();
+            fail("Should throw IllegalArgumentException for overflow-causing memory limit");
+        } catch (IllegalArgumentException e) {
+            assertTrue("Should throw IllegalArgumentException for overflow protection", true);
+        }
+    }
+
+    public void testMaxSafeMemoryLimit() throws Exception {
+        // Test that the maximum safe value works without throwing
+        long maxSafeMB = Long.MAX_VALUE / (1024L * 1024L);
+
+        try {
+            SpiceClient client = SpiceClient.builder()
+                    .withArrowMemoryLimitMB(maxSafeMB) // Maximum safe value
+                    .build();
+
+            assertTrue("Maximum safe memory limit should work", true);
+            client.close();
+        } catch (Exception e) {
+            fail("Should not throw exception: " + e.getMessage());
+        }
+    }
 }