From a39d9b3849cae3b6648c30c02fa39c866c3cfb01 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Mon, 18 Sep 2023 11:03:10 +0200
Subject: [PATCH] feat: support setting core pool size for async API in system
 property

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
---
 .../google/cloud/spanner/SpannerOptions.java  | 25 ++++++++++-
 .../cloud/spanner/SpannerOptionsTest.java     | 44 +++++++++++++++++++
 2 files changed, 68 insertions(+), 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
index 5df163ecd87..5d5498e3b68 100644
--- 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
index 7061e255758..924a36eb645 100644
--- 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,49 @@ 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(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));
+  }
+
+  static <V> V runWithSystemProperty(
+      String propertyName, String propertyValue, Callable<V> callable) throws Exception {
+    String currentValue = System.getProperty(propertyName);
+    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 =