feat: support setting core pool size for async API in system property (…

…#2632) Add support for setting the core pool size for the thread pool that is used for the Async API using a system property. This allows users to change this property without having to change code, and to set the property if they are using a framework that builds on top of the client library, but that does not have a configuration option for setting a custom executor provider. Fixes #2631
googleapis · Oct 4, 2023 · e51c55d · e51c55d
1 parent 7be4df1
commit e51c55d
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 1 deletion.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java
@@ -524,7 +524,30 @@ public static FixedCloseableExecutorProvider create(ScheduledExecutorService exe
    */
   @VisibleForTesting
   static CloseableExecutorProvider createDefaultAsyncExecutorProvider() {
-    return createAsyncExecutorProvider(8, 60L, TimeUnit.SECONDS);
+    return createAsyncExecutorProvider(
+        getDefaultAsyncExecutorProviderCoreThreadCount(), 60L, TimeUnit.SECONDS);
+  }
+
+  @VisibleForTesting
+  static int getDefaultAsyncExecutorProviderCoreThreadCount() {
+    String propertyName = "SPANNER_ASYNC_NUM_CORE_THREADS";
+    String propertyValue = System.getProperty(propertyName, "8");
+    try {
+      int corePoolSize = Integer.parseInt(propertyValue);
+      if (corePoolSize < 0) {
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT,
+            String.format(
+                "The value for %s must be >=0. Invalid value: %s", propertyName, propertyValue));
+      }
+      return corePoolSize;
+    } catch (NumberFormatException exception) {
+      throw SpannerExceptionFactory.newSpannerException(
+          ErrorCode.INVALID_ARGUMENT,
+          String.format(
+              "The %s system property must be a valid integer. The value %s could not be parsed as an integer.",
+              propertyName, propertyValue));
+    }
   }
 
   /**

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java
@@ -63,6 +63,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.Callable;
 import java.util.concurrent.ScheduledExecutorService;
 import javax.annotation.Nonnull;
 import org.junit.Test;
@@ -913,6 +914,62 @@ public void testCustomAsyncExecutorProvider() {
     assertSame(service, options.getAsyncExecutorProvider().getExecutor());
   }
 
+  @Test
+  public void testAsyncExecutorProviderCoreThreadCount() throws Exception {
+    assertEquals(8, SpannerOptions.getDefaultAsyncExecutorProviderCoreThreadCount());
+    String propertyName = "SPANNER_ASYNC_NUM_CORE_THREADS";
+    assertEquals(
+        Integer.valueOf(8),
+        runWithSystemProperty(
+            propertyName, null, SpannerOptions::getDefaultAsyncExecutorProviderCoreThreadCount));
+    assertEquals(
+        Integer.valueOf(16),
+        runWithSystemProperty(
+            propertyName, "16", SpannerOptions::getDefaultAsyncExecutorProviderCoreThreadCount));
+    assertEquals(
+        Integer.valueOf(1),
+        runWithSystemProperty(
+            propertyName, "1", SpannerOptions::getDefaultAsyncExecutorProviderCoreThreadCount));
+    assertThrows(
+        SpannerException.class,
+        () ->
+            runWithSystemProperty(
+                propertyName,
+                "foo",
+                SpannerOptions::getDefaultAsyncExecutorProviderCoreThreadCount));
+    assertThrows(
+        SpannerException.class,
+        () ->
+            runWithSystemProperty(
+                propertyName,
+                "-1",
+                SpannerOptions::getDefaultAsyncExecutorProviderCoreThreadCount));
+    assertThrows(
+        SpannerException.class,
+        () ->
+            runWithSystemProperty(
+                propertyName, "", SpannerOptions::getDefaultAsyncExecutorProviderCoreThreadCount));
+  }
+
+  static <V> V runWithSystemProperty(
+      String propertyName, String propertyValue, Callable<V> callable) throws Exception {
+    String currentValue = System.getProperty(propertyName);
+    if (propertyValue == null) {
+      System.clearProperty(propertyName);
+    } else {
+      System.setProperty(propertyName, propertyValue);
+    }
+    try {
+      return callable.call();
+    } finally {
+      if (currentValue == null) {
+        System.clearProperty(propertyName);
+      } else {
+        System.setProperty(propertyName, currentValue);
+      }
+    }
+  }
+
   @Test
   public void testDefaultNumChannelsWithGrpcGcpExtensionEnabled() {
     SpannerOptions options =