Fix executor metrics lifecycle

base-cluster/src/main/java/org/apache/bifromq/basecluster/AgentHost.java

@@ -26,7 +26,6 @@
 import com.google.common.base.Strings;
 import com.google.protobuf.ByteString;
 import io.micrometer.core.instrument.Metrics;
-import io.micrometer.core.instrument.binder.jvm.ExecutorServiceMetrics;
 import io.reactivex.rxjava3.core.Observable;
 import io.reactivex.rxjava3.core.Scheduler;
 import io.reactivex.rxjava3.disposables.CompositeDisposable;
@@ -36,9 +35,11 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.atomic.AtomicReference;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.bifromq.base.util.MonitoredExecutors;
 import org.apache.bifromq.basecluster.agent.proto.AgentEndpoint;
 import org.apache.bifromq.basecluster.fd.FailureDetector;
 import org.apache.bifromq.basecluster.fd.IFailureDetector;
@@ -69,6 +70,7 @@ final class AgentHost implements IAgentHost {
     private final ITransport transport;
     private final IMessenger messenger;
     private final IHostMemberList memberList;
+    private final ScheduledExecutorService hostExecutor;
     private final Scheduler hostScheduler;
     private final MemberSelector memberSelector;
     private final AutoHealer healer;
@@ -89,9 +91,11 @@ final class AgentHost implements IAgentHost {
             .maxHealthScore(options.awarenessMaxMultiplier())
             .retransmitMultiplier(options.retransmitMultiplier())
             .spreadPeriod(options.gossipPeriod());
-        hostScheduler = Schedulers.from(ExecutorServiceMetrics.monitor(Metrics.globalRegistry,
+        hostExecutor = MonitoredExecutors.monitor(Metrics.globalRegistry,
             new ScheduledThreadPoolExecutor(1, EnvProvider.INSTANCE.newThreadFactory("agent-host-scheduler", true)),
-            "agent-host-scheduler"));
+            "agent-host-scheduler[" + options.env() + "@" + options.addr() + ":"
+                + transport.bindAddress().getPort() + "]");
+        hostScheduler = Schedulers.from(hostExecutor);
         this.store = ICRDTStore.newInstance(options.crdtStoreOptions());
         this.messenger = Messenger.builder()
             .transport(transport)
@@ -197,6 +201,7 @@ public void close() {
                     memberSelector.stop();
                     disposables.dispose();
                     hostScheduler.shutdown();
+                    hostExecutor.shutdown();
                     state.set(State.SHUTDOWN);
                 }).join();
             log.debug("AgentHost stopped");

base-cluster/src/test/java/org/apache/bifromq/basecluster/AgentHostExecutorLifecycleTest.java

@@ -0,0 +1,125 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.bifromq.basecluster;
+
+import static org.awaitility.Awaitility.await;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotEquals;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertTrue;
+
+import io.micrometer.core.instrument.Meter;
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Timer;
+import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
+import java.time.Duration;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import org.testng.annotations.Test;
+
+public class AgentHostExecutorLifecycleTest {
+    @Test
+    public void closeStopsScheduler() throws Exception {
+        SimpleMeterRegistry registry = new SimpleMeterRegistry();
+        Metrics.addRegistry(registry);
+        Set<Thread> before = schedulerThreads();
+        try (IAgentHost host = IAgentHost.newInstance(options())) {
+            await().atMost(Duration.ofSeconds(10)).until(() -> schedulerThreads().size() > before.size());
+            Set<Thread> started = schedulerThreads();
+            started.removeAll(before);
+            assertEquals(started.size(), 1);
+            CompletableFuture.runAsync(host::close).get(15, TimeUnit.SECONDS);
+            for (Thread thread : started) {
+                thread.join(5000);
+                assertFalse(thread.isAlive(), thread.getName());
+            }
+            assertTrue(registry.getMeters().stream().map(Meter::getId)
+                .noneMatch(id -> id.getName().startsWith("executor")
+                    && id.getTag("name").startsWith("agent-host-scheduler")));
+        } finally {
+            Metrics.removeRegistry(registry);
+            registry.close();
+        }
+    }
+
+    @Test
+    public void hostMetricsAreIndependentAndCanBeRecreated() throws Exception {
+        SimpleMeterRegistry registry = new SimpleMeterRegistry();
+        Metrics.addRegistry(registry);
+        try (IAgentHost first = IAgentHost.newInstance(options());
+             IAgentHost second = IAgentHost.newInstance(options())) {
+            String firstName = executorName(first);
+            String secondName = executorName(second);
+            assertNotEquals(firstName, secondName);
+            assertEquals(executorMeters(registry, firstName).size(), 11);
+            assertEquals(executorMeters(registry, secondName).size(), 11);
+            for (Meter meter : executorMeters(registry, firstName)) {
+                assertEquals(meter.getId().getTags().size(), 1);
+                assertEquals(meter.getId().getTags().get(0).getKey(), "name");
+            }
+            Timer firstTimer = registry.get("executor").tag("name", firstName).timer();
+            Timer secondTimer = registry.get("executor").tag("name", secondName).timer();
+            await().atMost(Duration.ofSeconds(10)).until(() -> firstTimer.count() > 0 && secondTimer.count() > 0);
+            CompletableFuture.runAsync(first::close).get(15, TimeUnit.SECONDS);
+            assertTrue(executorMeters(registry, firstName).isEmpty());
+            assertEquals(executorMeters(registry, secondName).size(), 11);
+            long secondCount = secondTimer.count();
+            second.host("still-running");
+            await().atMost(Duration.ofSeconds(10)).until(() -> secondTimer.count() > secondCount);
+            try (IAgentHost replacement = IAgentHost.newInstance(options().port(first.local().getPort()))) {
+                assertEquals(executorName(replacement), firstName);
+                Timer newTimer = registry.get("executor").tag("name", firstName).timer();
+                assertNotSame(newTimer, firstTimer);
+                assertEquals(executorMeters(registry, firstName).size(), 11);
+                first.close();
+                assertEquals(executorMeters(registry, firstName).size(), 11);
+                await().atMost(Duration.ofSeconds(10)).until(() -> newTimer.count() > 0);
+            }
+            assertTrue(executorMeters(registry, firstName).isEmpty());
+        } finally {
+            Metrics.removeRegistry(registry);
+            registry.close();
+        }
+    }
+
+    private static AgentHostOptions options() {
+        return AgentHostOptions.builder().addr("127.0.0.1").env("executor-lifecycle")
+            .gossipPeriod(Duration.ofMillis(20)).build();
+    }
+
+    private static String executorName(IAgentHost host) {
+        return "agent-host-scheduler[" + host.env() + "@" + host.local().getAddress() + ":"
+            + host.local().getPort() + "]";
+    }
+
+    private static List<Meter> executorMeters(SimpleMeterRegistry registry, String name) {
+        return registry.getMeters().stream().filter(meter -> meter.getId().getName().startsWith("executor")
+            && name.equals(meter.getId().getTag("name"))).toList();
+    }
+
+    private static Set<Thread> schedulerThreads() {
+        return Thread.getAllStackTraces().keySet().stream().filter(thread -> thread.isAlive()
+            && thread.getName().startsWith("agent-host-scheduler")).collect(Collectors.toSet());
+    }
+}

base-kv/base-kv-local-engine-rocksdb/pom.xml

@@ -29,6 +29,10 @@
     <artifactId>base-kv-local-engine-rocksdb</artifactId>
 
     <dependencies>
+        <dependency>
+            <groupId>org.apache.bifromq</groupId>
+            <artifactId>base-util</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.apache.bifromq</groupId>
             <artifactId>base-kv-local-engine-spi</artifactId>

base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBCPableKVSpace.java

@@ -89,13 +89,17 @@ class RocksDBCPableKVSpace extends RocksDBKVSpace implements ICPableKVSpace {
         this.engine = engine;
         cpRootDir = new File(strVal(conf, DB_CHECKPOINT_ROOT_DIR), id);
         checkpoints = Caffeine.newBuilder().weakValues().build();
+        try {
+            Files.createDirectories(cpRootDir.getAbsoluteFile().toPath());
+            Files.createDirectories(spaceRootDir().getAbsoluteFile().toPath());
+            initOrMigrateCurrentDBDir();
+        } catch (IOException e) {
+            // The parent has already registered the compactor meters.
+            super.doClose();
+            throw e;
+        }
         writeOptions = new WriteOptions().setDisableWAL(true);
-        Files.createDirectories(cpRootDir.getAbsoluteFile().toPath());
         metricMgr = new MetricManager(tags);
-        // Ensure space root exists and initialize current generation directory.
-        // The legacy layout migration is done at constructor time to keep open() path clean.
-        Files.createDirectories(spaceRootDir().getAbsoluteFile().toPath());
-        initOrMigrateCurrentDBDir();
     }
 
     @Override
@@ -144,20 +148,22 @@ public IKVSpaceMigratableWriter toWriter() {
 
     @Override
     protected void doClose() {
-        logger.debug("Flush RocksDBCPableKVSpace[{}] before closing", id);
-        try (FlushOptions flushOptions = new FlushOptions().setWaitForFlush(true)) {
-            active.get().db.flush(flushOptions);
-        } catch (Throwable e) {
-            logger.error("Flush RocksDBCPableKVSpace[{}] error", id, e);
+        super.doClose();
+        RocksDBCPableKVSpaceEpochHandle h = active.get();
+        if (h != null) {
+            logger.debug("Flush RocksDBCPableKVSpace[{}] before closing", id);
+            try (FlushOptions flushOptions = new FlushOptions().setWaitForFlush(true)) {
+                h.db.flush(flushOptions);
+            } catch (Throwable e) {
+                logger.error("Flush RocksDBCPableKVSpace[{}] error", id, e);
+            }
         }
         metricMgr.close();
         checkpoints.asMap().forEach((cpId, cp) -> cp.close());
         writeOptions.close();
-        RocksDBCPableKVSpaceEpochHandle h = active.get();
         if (h != null) {
             h.close();
         }
-        super.doClose();
     }
 
     @Override
@@ -181,6 +187,7 @@ protected void doOpen() {
             loadLatestCheckpoint();
             super.doOpen();
         } catch (Throwable e) {
+            close();
             throw new KVEngineException("Failed to open CPable KVSpace", e);
         }
     }

base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBKVSpace.java

@@ -19,6 +19,7 @@
 
 package org.apache.bifromq.basekv.localengine.rocksdb;
 
+import static com.google.common.util.concurrent.Uninterruptibles.awaitTerminationUninterruptibly;
 import static org.apache.bifromq.basekv.localengine.StructUtil.boolVal;
 import static org.apache.bifromq.basekv.localengine.StructUtil.numVal;
 import static org.apache.bifromq.basekv.localengine.StructUtil.strVal;
@@ -39,7 +40,6 @@
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Tags;
 import io.micrometer.core.instrument.Timer;
-import io.micrometer.core.instrument.binder.jvm.ExecutorServiceMetrics;
 import java.io.File;
 import java.io.IOException;
 import java.util.Collections;
@@ -50,6 +50,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import lombok.SneakyThrows;
+import org.apache.bifromq.base.util.MonitoredExecutors;
 import org.apache.bifromq.baseenv.EnvProvider;
 import org.apache.bifromq.basekv.localengine.AbstractKVSpace;
 import org.apache.bifromq.basekv.localengine.ISyncContext;
@@ -89,10 +90,6 @@ public RocksDBKVSpace(String id,
         syncContext = new SyncContext();
         metadataRefresher = syncContext.refresher();
         compacting = new AtomicBoolean(false);
-        compactionExecutor = ExecutorServiceMetrics.monitor(Metrics.globalRegistry, new ThreadPoolExecutor(1, 1,
-                0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>(),
-                EnvProvider.INSTANCE.newThreadFactory("kvspace-compactor-" + id)),
-            "compactor", "kvspace", Tags.of(tags));
         if (boolVal(conf, MANUAL_COMPACTION)) {
             int minKeys = (int) numVal(conf, COMPACT_MIN_TOMBSTONE_KEYS);
             int minRanges = (int) numVal(conf, COMPACT_MIN_TOMBSTONE_RANGES);
@@ -103,6 +100,10 @@ public RocksDBKVSpace(String id,
             this.writeStats = NoopWriteStatsRecorder.INSTANCE;
         }
         keySpaceDBDir = new File(strVal(conf, DB_ROOT_DIR), id);
+        compactionExecutor = MonitoredExecutors.monitor(Metrics.globalRegistry, new ThreadPoolExecutor(1, 1,
+                0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>(),
+                EnvProvider.INSTANCE.newThreadFactory("kvspace-compactor-" + id)),
+            "compactor", "kvspace", this.tags);
     }
 
     @Override
@@ -130,6 +131,8 @@ protected void publishMetadata(Map<ByteString, ByteString> metadataUpdates) {
     @Override
     protected void doClose() {
         logger.debug("Close key range[{}]", id);
+        compactionExecutor.shutdown();
+        awaitTerminationUninterruptibly(compactionExecutor);
         if (spaceMetrics != null) {
             spaceMetrics.close();
         }

base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBWALableKVSpace.java

@@ -19,6 +19,7 @@
 
 package org.apache.bifromq.basekv.localengine.rocksdb;
 
+import static com.google.common.util.concurrent.Uninterruptibles.awaitTerminationUninterruptibly;
 import static org.apache.bifromq.basekv.localengine.StructUtil.boolVal;
 import static org.apache.bifromq.basekv.localengine.rocksdb.RocksDBDefaultConfigs.ASYNC_WAL_FLUSH;
 import static org.apache.bifromq.basekv.localengine.rocksdb.RocksDBDefaultConfigs.FSYNC_WAL;
@@ -27,15 +28,14 @@
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Tags;
 import io.micrometer.core.instrument.Timer;
-import io.micrometer.core.instrument.binder.jvm.ExecutorServiceMetrics;
 import java.nio.file.Files;
-import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import org.apache.bifromq.base.util.MonitoredExecutors;
 import org.apache.bifromq.baseenv.EnvProvider;
 import org.apache.bifromq.basekv.localengine.IKVSpaceRefreshableReader;
 import org.apache.bifromq.basekv.localengine.IKVSpaceWriter;
@@ -66,12 +66,11 @@ class RocksDBWALableKVSpace extends RocksDBKVSpace implements IWALableKVSpace {
         if (isSyncWALFlush()) {
             writeOptions.setSync(isFsyncWAL());
         }
-        flushExecutor = ExecutorServiceMetrics.monitor(Metrics.globalRegistry, new ThreadPoolExecutor(1, 1,
-                0L, TimeUnit.MILLISECONDS,
-                new LinkedBlockingQueue<>(),
-                EnvProvider.INSTANCE.newThreadFactory("kvspace-flusher-" + id)), "flusher", "kvspace",
-            Tags.of(tags));
         metricMgr = new MetricManager(tags);
+        flushExecutor = MonitoredExecutors.monitor(Metrics.globalRegistry, new ThreadPoolExecutor(1, 1,
+                0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>(),
+                EnvProvider.INSTANCE.newThreadFactory("kvspace-flusher-" + id)),
+            "flusher", "kvspace", this.tags);
     }
 
     @Override
@@ -82,6 +81,7 @@ protected void doOpen() {
             handle = new RocksDBWALableKVSpaceEpochHandle(id, spaceRootDir(), this.conf, logger, tags);
             super.doOpen();
         } catch (Throwable e) {
+            close();
             throw new KVEngineException("Failed to open WALable KVSpace", e);
         }
     }
@@ -93,22 +93,15 @@ protected RocksDBWALableKVSpaceEpochHandle handle() {
 
     @Override
     protected void doClose() {
-        final CompletableFuture<Long> flushTaskFuture = Optional.ofNullable(flushFutureRef.get()).orElseGet(() -> {
-            CompletableFuture<Long> lastOne = new CompletableFuture<>();
-            flushExecutor.submit(() -> lastOne.complete(System.nanoTime()));
-            return lastOne;
-        });
+        super.doClose();
         flushExecutor.shutdown();
-        try {
-            flushTaskFuture.join();
-        } catch (Throwable e) {
-            logger.debug("Flush error during closing", e);
-        }
+        awaitTerminationUninterruptibly(flushExecutor);
         writeOptions.close();
         metricMgr.close();
         // close handle
-        handle.close();
-        super.doClose();
+        if (handle != null) {
+            handle.close();
+        }
     }
 
     @Override