Allow using more threads to refresh cassandra pools

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/cassandra/CassandraKeyValueServiceConfig.java

@@ -381,4 +381,13 @@ default void check() {
                 localHostWeighting() >= 0.0 && localHostWeighting() <= 1.0,
                 "'localHostWeighting' must be between 0 and 1 inclusive");
     }
+
+    /**
+     * Number of threads to be used across ALL transaction managers to refresh the client pool. It is suggested to use
+     * one thread per approximately 10 transaction managers that are expected to be used.
+     */
+    @Value.Default
+    default int numPoolRefreshingThreads() {
+        return 1;
+    }
 }

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraClientPoolImpl.java

@@ -32,8 +32,8 @@
 import com.palantir.atlasdb.keyvalue.cassandra.pool.CassandraService;
 import com.palantir.atlasdb.util.MetricsManager;
 import com.palantir.common.base.FunctionCheckedException;
+import com.palantir.common.concurrent.InitializeableScheduledExecutorServiceSupplier;
 import com.palantir.common.concurrent.NamedThreadFactory;
-import com.palantir.common.concurrent.PTExecutors;
 import com.palantir.logsafe.SafeArg;
 import com.palantir.logsafe.UnsafeArg;
 import com.palantir.logsafe.exceptions.SafeIllegalStateException;
@@ -73,8 +73,9 @@
  **/
 @SuppressWarnings("checkstyle:FinalClass") // non-final for mocking
 public class CassandraClientPoolImpl implements CassandraClientPool {
-    private static final ScheduledExecutorService sharedRefreshDaemon =
-            PTExecutors.newScheduledThreadPool(1, new NamedThreadFactory("CassandraClientPoolRefresh", true));
+    private static final InitializeableScheduledExecutorServiceSupplier SHARED_EXECUTOR_SUPPLIER =
+            new InitializeableScheduledExecutorServiceSupplier(
+                    new NamedThreadFactory("CassandraClientPoolRefresh", true));
 
     private class InitializingWrapper extends AsyncInitializer implements AutoDelegate_CassandraClientPool {
         @Override
@@ -141,14 +142,14 @@ static CassandraClientPoolImpl createImplForTest(
             MetricsManager metricsManager,
             CassandraKeyValueServiceConfig config,
             StartupChecks startupChecks,
-            ScheduledExecutorService refreshDaemon,
+            InitializeableScheduledExecutorServiceSupplier initializeableExecutorSupplier,
             Blacklist blacklist,
             CassandraService cassandra) {
         CassandraRequestExceptionHandler exceptionHandler = testExceptionHandler(blacklist);
         CassandraClientPoolImpl cassandraClientPool = new CassandraClientPoolImpl(
                 config,
                 startupChecks,
-                refreshDaemon,
+                initializeableExecutorSupplier,
                 exceptionHandler,
                 blacklist,
                 cassandra,
@@ -189,7 +190,7 @@ private CassandraClientPoolImpl(
         this(
                 config,
                 startupChecks,
-                sharedRefreshDaemon,
+                SHARED_EXECUTOR_SUPPLIER,
                 exceptionHandler,
                 blacklist,
                 new CassandraService(metricsManager, config, blacklist, metrics),
@@ -199,14 +200,15 @@ private CassandraClientPoolImpl(
     private CassandraClientPoolImpl(
             CassandraKeyValueServiceConfig config,
             StartupChecks startupChecks,
-            ScheduledExecutorService refreshDaemon,
+            InitializeableScheduledExecutorServiceSupplier initializeableExecutorSupplier,
             CassandraRequestExceptionHandler exceptionHandler,
             Blacklist blacklist,
             CassandraService cassandra,
             CassandraClientPoolMetrics metrics) {
         this.config = config;
         this.startupChecks = startupChecks;
-        this.refreshDaemon = refreshDaemon;
+        initializeableExecutorSupplier.initialize(config.numPoolRefreshingThreads());
+        this.refreshDaemon = initializeableExecutorSupplier.get();
         this.blacklist = blacklist;
         this.exceptionHandler = exceptionHandler;
         this.cassandra = cassandra;

atlasdb-cassandra/src/test/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraClientPoolTest.java

@@ -31,6 +31,7 @@
 import com.palantir.atlasdb.keyvalue.cassandra.pool.CassandraService;
 import com.palantir.atlasdb.util.MetricsManagers;
 import com.palantir.common.base.FunctionCheckedException;
+import com.palantir.common.concurrent.InitializeableScheduledExecutorServiceSupplier;
 import com.palantir.common.exception.AtlasDbDependencyException;
 import com.palantir.conjure.java.api.config.service.HumanReadableDuration;
 import com.palantir.tritium.metrics.registry.DefaultTaggedMetricRegistry;
@@ -378,7 +379,7 @@ private CassandraClientPoolImpl createClientPool() {
                 MetricsManagers.createForTests(),
                 config,
                 CassandraClientPoolImpl.StartupChecks.DO_NOT_RUN,
-                deterministicExecutor,
+                InitializeableScheduledExecutorServiceSupplier.createForTests(deterministicExecutor),
                 blacklist,
                 cassandra);
     }

changelog/@unreleased/pr-5900.v2.yml

@@ -0,0 +1,6 @@
+type: improvement
+improvement:
+  description: Cassandra users can now specify the number of threads to be used for
+    refreshing pools
+  links:
+  - https://github.com/palantir/atlasdb/pull/5900

commons-executors/src/main/java/com/palantir/common/concurrent/InitializeableScheduledExecutorServiceSupplier.java

@@ -0,0 +1,56 @@
+/*
+ * (c) Copyright 2022 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.palantir.common.concurrent;
+
+import com.palantir.logsafe.Preconditions;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+
+/**
+ * A supplier of ScheduledExecutorService that can be initialized with the desired number of threads at runtime. Must
+ * be initialized before it is first used, and once initialized, further calls to the {@link #initialize(int)} method
+ * are ignored.
+ */
+public final class InitializeableScheduledExecutorServiceSupplier implements Supplier<ScheduledExecutorService> {
+    private final ThreadFactory threadFactory;
+    private final AtomicReference<ScheduledExecutorService> delegate = new AtomicReference<>();
+
+    public InitializeableScheduledExecutorServiceSupplier(ThreadFactory threadFactory) {
+        this.threadFactory = threadFactory;
+    }
+
+    private InitializeableScheduledExecutorServiceSupplier(ScheduledExecutorService executor, ThreadFactory factory) {
+        threadFactory = factory;
+        delegate.set(executor);
+    }
+
+    public static InitializeableScheduledExecutorServiceSupplier createForTests(ScheduledExecutorService executor) {
+        return new InitializeableScheduledExecutorServiceSupplier(executor, new NamedThreadFactory("fake", true));
+    }
+
+    public void initialize(int numThreads) {
+        delegate.compareAndSet(null, PTExecutors.newScheduledThreadPool(numThreads, threadFactory));
+    }
+
+    @Override
+    public ScheduledExecutorService get() {
+        Preconditions.checkState(delegate.get() != null, "Executor must be initialized");
+        return delegate.get();
+    }
+}

commons-executors/src/test/java/com/palantir/common/concurrent/InitializeableScheduledExecutorServiceSupplierTest.java

@@ -0,0 +1,65 @@
+/*
+ * (c) Copyright 2022 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.palantir.common.concurrent;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.junit.Test;
+
+public class InitializeableScheduledExecutorServiceSupplierTest {
+    private final InitializeableScheduledExecutorServiceSupplier supplier =
+            new InitializeableScheduledExecutorServiceSupplier(new NamedThreadFactory("test", true));
+
+    @Test
+    public void uninitializedSupplierThrows() {
+        assertThatThrownBy(supplier::get).isInstanceOf(IllegalStateException.class);
+    }
+
+    @Test
+    public void canInitializeSupplier() throws ExecutionException, InterruptedException {
+        supplier.initialize(1);
+
+        assertThat(supplier.get().schedule(() -> 1L, 0, TimeUnit.SECONDS).get()).isEqualTo(1L);
+        supplier.get().shutdownNow();
+    }
+
+    @Test
+    public void originalNumberOfThreadsIsRespected() throws InterruptedException {
+        InitializeableScheduledExecutorServiceSupplier supplier =
+                new InitializeableScheduledExecutorServiceSupplier(new NamedThreadFactory("test", true));
+        supplier.initialize(1);
+        supplier.initialize(2);
+
+        CountDownLatch latch = new CountDownLatch(2);
+        Runnable task = () -> {
+            latch.countDown();
+            try {
+                latch.await();
+            } catch (InterruptedException e) {
+                // ok cool
+            }
+        };
+
+        supplier.get().schedule(task, 0, TimeUnit.SECONDS);
+        supplier.get().schedule(task, 0, TimeUnit.SECONDS);
+        assertThat(latch.await(1, TimeUnit.SECONDS)).isFalse();
+    }
+}