Fix UDP socket leak in CentralDogma-to-CentralDogma mirroring (#1328)

Motivation:

CentralDogma-to-CentralDogma mirrors were creating a new
ArmeriaCentralDogmaClient on every scheduler tick, allocating a fresh
Armeria WebClient, connection pool, and UDP socket each time. These
sockets were not reclaimed promptly, leading to too many open
files.

Modifications:

- Add CentralDogma.withAccessToken(String) as a default interface method
  that returns a derived client sharing the same underlying WebClient.
  Calling close() on the derived client is a no-op; the base client owns
  the connection resources.
- Add Mirror.setBaseClientPool(ConcurrentHashMap<String, Object>) as a
  default no-op method so MirrorSchedulingService can inject the pool
  without coupling the server module to client types.

Result:

- Mirrors pointing at the same remote host share one Armeria WebClient,
  eliminating per-run DNS/TLS/connection overhead and stopping UDP socket
  accumulation.

Author: minwoox

client/java-armeria/src/main/java/com/linecorp/centraldogma/internal/client/armeria/ArmeriaCentralDogma.java

@@ -77,6 +77,7 @@
 import com.linecorp.armeria.common.util.SafeCloseable;
 import com.linecorp.armeria.common.util.TimeoutMode;
 import com.linecorp.centraldogma.client.AbstractCentralDogma;
+import com.linecorp.centraldogma.client.CentralDogma;
 import com.linecorp.centraldogma.client.CentralDogmaRepository;
 import com.linecorp.centraldogma.client.RepositoryInfo;
 import com.linecorp.centraldogma.common.ApiRequestTimeoutException;
@@ -176,6 +177,14 @@ public ArmeriaCentralDogma(ScheduledExecutorService blockingTaskExecutor,
         this.whenReady = whenReady;
     }
 
+    @Override
+    public CentralDogma withAccessToken(String accessToken) {
+        requireNonNull(accessToken, "accessToken");
+        // Pass a no-op SafeCloseable: the base client owns the connection resources.
+        return new ArmeriaCentralDogma(executor(), client, accessToken,
+                                       () -> {}, meterRegistry(), null);
+    }
+
     @Override
     public CompletableFuture<Void> whenEndpointReady() {
         if (whenReady != null) {

client/java/src/main/java/com/linecorp/centraldogma/client/AbstractCentralDogma.java

@@ -74,6 +74,14 @@ protected final ScheduledExecutorService executor() {
         return blockingTaskExecutor;
     }
 
+    /**
+     * Returns the {@link MeterRegistry}, or {@code null} if none was configured.
+     */
+    @Nullable
+    protected final MeterRegistry meterRegistry() {
+        return meterRegistry;
+    }
+
     @Override
     public CentralDogmaRepository forRepo(String projectName, String repositoryName) {
         requireNonNull(projectName, "projectName");

client/java/src/main/java/com/linecorp/centraldogma/client/CentralDogma.java

@@ -53,6 +53,17 @@
  */
 public interface CentralDogma extends AutoCloseable {
 
+    /**
+     * Returns a derived {@link CentralDogma} that reuses the underlying connection pool of this client
+     * but authenticates with the given {@code accessToken}. Calling {@link #close()} on the derived client
+     * is a no-op; the connection pool is owned by this (base) client.
+     *
+     * @throws UnsupportedOperationException if this implementation does not support derived clients
+     */
+    default CentralDogma withAccessToken(String accessToken) {
+        throw new UnsupportedOperationException(getClass().getName() + " does not support withAccessToken");
+    }
+
     /**
      * Returns a new {@link CentralDogmaRepository} that is used to send a request to the specified
      * {@code projectName} and {@code repositoryName}.

client/java/src/main/java/com/linecorp/centraldogma/internal/client/ReplicationLagTolerantCentralDogma.java

@@ -106,6 +106,14 @@ public ReplicationLagTolerantCentralDogma(ScheduledExecutorService blockingTaskE
         this.currentReplicaHintSupplier = currentReplicaHintSupplier;
     }
 
+    @Override
+    public CentralDogma withAccessToken(String accessToken) {
+        requireNonNull(accessToken, "accessToken");
+        return new ReplicationLagTolerantCentralDogma(
+                executor(), delegate.withAccessToken(accessToken),
+                maxRetries, retryIntervalMillis, currentReplicaHintSupplier, meterRegistry());
+    }
+
     @Override
     public CompletableFuture<Void> createProject(String projectName) {
         return delegate.createProject(projectName);

it/mirror-listener/src/test/java/com/linecorp/centraldogma/it/mirror/listener/CustomMirrorListenerTest.java

@@ -127,8 +127,7 @@ protected MirrorResult mirrorRemoteToLocal(File workDir, CommandExecutor executo
         final DefaultMirrorAccessController ac = new DefaultMirrorAccessController();
         ac.setRepository(repositoryExtension.crudRepository());
         final MirrorSchedulingService service = new MirrorSchedulingService(
-                temporaryFolder, pm, new SimpleMeterRegistry(), 1, 1, 1, null, false,
-                ac);
+                temporaryFolder, pm, new SimpleMeterRegistry(), 1, 1, 1, null, ac);
         final CommandExecutor executor = mock(CommandExecutor.class);
         service.start(executor);
 

server-mirror-dogma/src/main/java/com/linecorp/centraldogma/server/internal/mirror/CentralDogmaMirror.java

@@ -29,6 +29,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 import org.jspecify.annotations.Nullable;
 import org.slf4j.Logger;
@@ -52,6 +53,7 @@
 import com.linecorp.centraldogma.common.PathPattern;
 import com.linecorp.centraldogma.common.RedundantChangeException;
 import com.linecorp.centraldogma.common.Revision;
+import com.linecorp.centraldogma.internal.CsrfToken;
 import com.linecorp.centraldogma.internal.Jackson;
 import com.linecorp.centraldogma.internal.Util;
 import com.linecorp.centraldogma.server.command.Command;
@@ -72,6 +74,10 @@ public final class CentralDogmaMirror extends AbstractMirror {
     private final String remoteProject;
     private final String remoteRepo;
 
+    // Injected by MirrorSchedulingService before each run; null only in tests / direct usage.
+    @Nullable
+    private volatile ConcurrentHashMap<String, Object> baseClientPool;
+
     public CentralDogmaMirror(String id, boolean enabled, Cron schedule, MirrorDirection direction,
                               Credential credential, Repository localRepo, String localPath,
                               RepositoryUri remoteUri, String remoteProject, String remoteRepo,
@@ -83,6 +89,11 @@ public CentralDogmaMirror(String id, boolean enabled, Cron schedule, MirrorDirec
         this.remoteRepo = requireNonNull(remoteRepo, "remoteRepo");
     }
 
+    @Override
+    public void setBaseClientPool(ConcurrentHashMap<String, Object> pool) {
+        baseClientPool = pool;
+    }
+
     @VisibleForTesting
     String remoteProject() {
         return remoteProject;
@@ -93,25 +104,64 @@ String remoteRepo() {
         return remoteRepo;
     }
 
-    private CentralDogma createRemoteClient(long maxNumBytes) throws UnknownHostException {
+    private static String baseClientPoolKey(URI uri) {
+        return uri.getHost() + ':' + uri.getPort() + ':' +
+               SCHEME_DOGMA_HTTPS.equals(uri.getScheme());
+    }
+
+    // maxNumBytes comes from MirrorSchedulingService.maxNumBytesPerMirror, which is fixed at server startup.
+    private synchronized CentralDogma createRemoteClient(long maxNumBytes) throws UnknownHostException {
         final URI uri = remoteUri().uri();
+        final ConcurrentHashMap<String, Object> pool = baseClientPool;
+
+        final Credential cred = credential();
+        final String token = cred instanceof AccessTokenCredential ?
+                             ((AccessTokenCredential) cred).accessToken() : CsrfToken.ANONYMOUS;
+
+        CentralDogma base;
+        if (pool != null) {
+            // Reuse or create a shared base client (MirrorSchedulingService owns the pool lifecycle).
+            final String key = baseClientPoolKey(uri);
+            final Object existing = pool.get(key);
+            if (existing instanceof CentralDogma) {
+                base = (CentralDogma) existing;
+            } else {
+                base = createClient(uri, maxNumBytes, null);
+                final Object prev = pool.putIfAbsent(key, base);
+                if (prev instanceof CentralDogma) {
+                    // Lost the race; close ours and reuse the winner.
+                    try {
+                        base.close();
+                    } catch (Exception e) {
+                        logger.warn("Failed to close the redundant base CentralDogma client: {}", uri, e);
+                    }
+                    base = (CentralDogma) prev;
+                }
+            }
+
+            // TODO(minwoox) Support mTLS authentication as well.
+            return base.withAccessToken(token);
+        }
+        // No pool injected (tests or direct usage without the scheduler).
+        return createClient(uri, maxNumBytes, token);
+    }
+
+    private static CentralDogma createClient(URI uri, long maxResponseLength, @Nullable String token)
+            throws UnknownHostException {
         final ArmeriaCentralDogmaBuilder builder = new ArmeriaCentralDogmaBuilder();
-        final int port = uri.getPort();
-        if (port > 0) {
-            builder.host(uri.getHost(), port);
+        if (uri.getPort() > 0) {
+            builder.host(uri.getHost(), uri.getPort());
         } else {
             builder.host(uri.getHost());
         }
-        if (SCHEME_DOGMA_HTTPS.equals(uri.getScheme())) {
-            builder.useTls(true);
-        }
-        final Credential cred = credential();
-        if (cred instanceof AccessTokenCredential) {
-            builder.accessToken(((AccessTokenCredential) cred).accessToken());
+        builder.useTls(SCHEME_DOGMA_HTTPS.equals(uri.getScheme()));
+        builder.clientConfigurator(cb -> cb.maxResponseLength(maxResponseLength));
+        // Mirrors run on a fixed schedule; a failed run simply retries on the next tick.
+        // Persistent health-check traffic is unnecessary and wasteful at scale.
+        builder.healthCheckIntervalMillis(0);
+        if (token != null) {
+            builder.accessToken(token);
         }
-        // TODO(minwoox) Support mTLS authentication as well.
-
-        builder.clientConfigurator(cb -> cb.maxResponseLength(maxNumBytes));
         return builder.build();
     }
 

server-mirror-git/src/test/java/com/linecorp/centraldogma/server/internal/mirror/MirrorSchedulingServiceTest.java

@@ -102,7 +102,7 @@ protected MirrorResult mirrorRemoteToLocal(File workDir, CommandExecutor executo
         when(mr.mirrors()).thenReturn(CompletableFuture.completedFuture(ImmutableList.of(mirror)));
 
         final MirrorSchedulingService service = new MirrorSchedulingService(
-                temporaryFolder, pm, new SimpleMeterRegistry(), 1, 1, 1, null, false,
+                temporaryFolder, pm, new SimpleMeterRegistry(), 1, 1, 1, null,
                 AlwaysAllowedMirrorAccessController.INSTANCE);
         final CommandExecutor executor = mock(CommandExecutor.class);
         service.start(executor);

server/src/main/java/com/linecorp/centraldogma/server/internal/mirror/DefaultMirroringServicePlugin.java

@@ -77,7 +77,6 @@ public synchronized CompletionStage<Void> start(PluginContext context) {
             final int maxNumFilesPerMirror;
             final long maxNumBytesPerMirror;
             final ZoneConfig zoneConfig;
-            final boolean runMigration;
 
             if (mirroringServicePluginConfig != null) {
                 numThreads = mirroringServicePluginConfig.numMirroringThreads();
@@ -93,13 +92,11 @@ public synchronized CompletionStage<Void> start(PluginContext context) {
                     logger.warn("No 'trustedHostKeys' configured in the mirroring service plugin config. " +
                                 "SSH mirror connections will accept any host key without verification.");
                 }
-                runMigration = mirroringServicePluginConfig.runMigration();
             } else {
                 numThreads = MirroringServicePluginConfig.INSTANCE.numMirroringThreads();
                 maxNumFilesPerMirror = MirroringServicePluginConfig.INSTANCE.maxNumFilesPerMirror();
                 maxNumBytesPerMirror = MirroringServicePluginConfig.INSTANCE.maxNumBytesPerMirror();
                 zoneConfig = null;
-                runMigration = true;
                 logger.warn("No 'trustedHostKeys' configured in the mirroring service plugin config. " +
                             "SSH mirror connections will accept any host key without verification.");
             }
@@ -108,7 +105,7 @@ public synchronized CompletionStage<Void> start(PluginContext context) {
                                                            context.meterRegistry(),
                                                            numThreads,
                                                            maxNumFilesPerMirror,
-                                                           maxNumBytesPerMirror, zoneConfig, runMigration,
+                                                           maxNumBytesPerMirror, zoneConfig,
                                                            context.mirrorAccessController());
             this.mirroringService = mirroringService;
         }

server/src/main/java/com/linecorp/centraldogma/server/internal/mirror/MirrorSchedulingService.java

@@ -27,6 +27,7 @@
 import java.util.List;
 import java.util.ServiceLoader;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -104,27 +105,28 @@ public static MirrorListener mirrorListener() {
     private final String currentZone;
     private final MirrorAccessController mirrorAccessController;
     private final AtomicInteger numActiveMirrors = new AtomicInteger();
+    // Shared base-client pool for CentralDogmaMirror instances. Keyed by "host:port:tls:maxBytes".
+    // Owned here so stop() can close all entries; injected into each mirror before it runs.
+    private final ConcurrentHashMap<String, Object> baseClientPool = new ConcurrentHashMap<>();
 
     private volatile CommandExecutor commandExecutor;
     private volatile ListeningScheduledExecutorService scheduler;
     private volatile ListeningExecutorService worker;
     private volatile boolean closing;
 
+    @Nullable
     private ZonedDateTime lastExecutionTime;
     private final MeterRegistry meterRegistry;
-    // Used to disable in the tests.
-    private final boolean runMigration;
 
     @VisibleForTesting
     public MirrorSchedulingService(File workDir, ProjectManager projectManager, MeterRegistry meterRegistry,
                                    int numThreads, int maxNumFilesPerMirror, long maxNumBytesPerMirror,
-                                   @Nullable ZoneConfig zoneConfig, boolean runMigration,
+                                   @Nullable ZoneConfig zoneConfig,
                                    MirrorAccessController mirrorAccessController) {
 
         this.workDir = requireNonNull(workDir, "workDir");
         this.projectManager = requireNonNull(projectManager, "projectManager");
         this.meterRegistry = requireNonNull(meterRegistry, "meterRegistry");
-        this.runMigration = runMigration;
 
         checkArgument(numThreads > 0, "numThreads: %s (expected: > 0)", numThreads);
         checkArgument(maxNumFilesPerMirror > 0,
@@ -227,6 +229,16 @@ public synchronized void stop() {
         } finally {
             this.scheduler = null;
             this.worker = null;
+            baseClientPool.forEach((key, client) -> {
+                if (client instanceof AutoCloseable) {
+                    try {
+                        ((AutoCloseable) client).close();
+                    } catch (Exception e) {
+                        logger.warn("Failed to close base CentralDogma client for key: {}", key, e);
+                    }
+                }
+            });
+            baseClientPool.clear();
         }
     }
 
@@ -310,6 +322,7 @@ private void scheduleMirrors() {
                               }
                               try {
                                   if (m.nextExecutionTime(currentLastExecutionTime).compareTo(now) < 0) {
+                                      m.setBaseClientPool(baseClientPool);
                                       runAsync(new MirrorTask(m, User.SYSTEM, Instant.now(),
                                                               currentZone, true));
                                   }
@@ -334,7 +347,10 @@ public CompletableFuture<Void> mirror() {
                     }
                     try {
                         p.metaRepo().mirrors().get(5, TimeUnit.SECONDS)
-                         .forEach(m -> run(new MirrorTask(m, User.SYSTEM, Instant.now(), currentZone, false)));
+                         .forEach(m -> {
+                             m.setBaseClientPool(baseClientPool);
+                             run(new MirrorTask(m, User.SYSTEM, Instant.now(), currentZone, false));
+                         });
                     } catch (InterruptedException | TimeoutException | ExecutionException e) {
                         throw new IllegalStateException(
                                 "Failed to load mirror list with in 5 seconds. project: " + p.name(), e);

server/src/main/java/com/linecorp/centraldogma/server/mirror/Mirror.java

@@ -19,6 +19,7 @@
 import java.net.URI;
 import java.time.Instant;
 import java.time.ZonedDateTime;
+import java.util.concurrent.ConcurrentHashMap;
 
 import org.jspecify.annotations.Nullable;
 
@@ -131,4 +132,11 @@ default String remoteBranch() {
      */
     MirrorResult mirror(File workDir, CommandExecutor executor, int maxNumFiles, long maxNumBytes,
                         Instant triggeredTime);
+
+    /**
+     * Injects the shared base-client pool managed by the mirroring scheduler. The pool maps a string key
+     * (encoding host, port, TLS flag, and maxResponseLength) to a shared base client instance. Only
+     * CentralDogma-type mirrors use this; the default implementation is a no-op.
+     */
+    default void setBaseClientPool(ConcurrentHashMap<String, Object> baseClientPool) {}
 }