diff --git a/base-cluster/src/main/java/org/apache/bifromq/basecluster/AgentHost.java b/base-cluster/src/main/java/org/apache/bifromq/basecluster/AgentHost.java index 61e8f7209..e4fb29350 100644 --- a/base-cluster/src/main/java/org/apache/bifromq/basecluster/AgentHost.java +++ b/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"); diff --git a/base-cluster/src/test/java/org/apache/bifromq/basecluster/AgentHostExecutorLifecycleTest.java b/base-cluster/src/test/java/org/apache/bifromq/basecluster/AgentHostExecutorLifecycleTest.java new file mode 100644 index 000000000..7ca4059e6 --- /dev/null +++ b/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 before = schedulerThreads(); + try (IAgentHost host = IAgentHost.newInstance(options())) { + await().atMost(Duration.ofSeconds(10)).until(() -> schedulerThreads().size() > before.size()); + Set 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 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 schedulerThreads() { + return Thread.getAllStackTraces().keySet().stream().filter(thread -> thread.isAlive() + && thread.getName().startsWith("agent-host-scheduler")).collect(Collectors.toSet()); + } +} diff --git a/base-kv/base-kv-local-engine-rocksdb/pom.xml b/base-kv/base-kv-local-engine-rocksdb/pom.xml index 81ea28921..7efdff7a5 100644 --- a/base-kv/base-kv-local-engine-rocksdb/pom.xml +++ b/base-kv/base-kv-local-engine-rocksdb/pom.xml @@ -29,6 +29,10 @@ base-kv-local-engine-rocksdb + + org.apache.bifromq + base-util + org.apache.bifromq base-kv-local-engine-spi diff --git a/base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBCPableKVSpace.java b/base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBCPableKVSpace.java index 665fe8f52..b3ce760f7 100644 --- a/base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBCPableKVSpace.java +++ b/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); } } diff --git a/base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBKVSpace.java b/base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBKVSpace.java index adae463e9..90bee6c06 100644 --- a/base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBKVSpace.java +++ b/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 metadataUpdates) { @Override protected void doClose() { logger.debug("Close key range[{}]", id); + compactionExecutor.shutdown(); + awaitTerminationUninterruptibly(compactionExecutor); if (spaceMetrics != null) { spaceMetrics.close(); } diff --git a/base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBWALableKVSpace.java b/base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBWALableKVSpace.java index 11e0fd192..c0cc4273c 100644 --- a/base-kv/base-kv-local-engine-rocksdb/src/main/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBWALableKVSpace.java +++ b/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 flushTaskFuture = Optional.ofNullable(flushFutureRef.get()).orElseGet(() -> { - CompletableFuture 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 diff --git a/base-kv/base-kv-local-engine-rocksdb/src/test/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBExecutorLifecycleTest.java b/base-kv/base-kv-local-engine-rocksdb/src/test/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBExecutorLifecycleTest.java new file mode 100644 index 000000000..e3f6ada9a --- /dev/null +++ b/base-kv/base-kv-local-engine-rocksdb/src/test/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBExecutorLifecycleTest.java @@ -0,0 +1,336 @@ +/* + * 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.basekv.localengine.rocksdb; + +import static org.apache.bifromq.basekv.localengine.StructUtil.toValue; +import static org.apache.bifromq.basekv.localengine.rocksdb.RocksDBDefaultConfigs.COMPACT_MIN_TOMBSTONE_KEYS; +import static org.apache.bifromq.basekv.localengine.rocksdb.RocksDBDefaultConfigs.DB_CHECKPOINT_ROOT_DIR; +import static org.apache.bifromq.basekv.localengine.rocksdb.RocksDBDefaultConfigs.DB_ROOT_DIR; +import static org.apache.bifromq.basekv.localengine.rocksdb.RocksDBHelper.deleteDir; +import static org.awaitility.Awaitility.await; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.spy; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotSame; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + +import com.google.protobuf.ByteString; +import com.google.protobuf.Struct; +import io.micrometer.core.instrument.Gauge; +import io.micrometer.core.instrument.Meter; +import io.micrometer.core.instrument.Metrics; +import io.micrometer.core.instrument.Timer; +import io.micrometer.core.instrument.distribution.DistributionStatisticConfig; +import io.micrometer.core.instrument.distribution.pause.PauseDetector; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.apache.bifromq.basekv.localengine.ICPableKVSpace; +import org.apache.bifromq.basekv.localengine.IKVSpace; +import org.apache.bifromq.basekv.localengine.IKVSpaceReader; +import org.apache.bifromq.basekv.localengine.IKVSpaceWriter; +import org.apache.bifromq.basekv.localengine.IWALableKVSpace; +import org.apache.bifromq.basekv.localengine.KVEngineException; +import org.rocksdb.Options; +import org.rocksdb.RocksDB; +import org.rocksdb.RocksDBException; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +public class RocksDBExecutorLifecycleTest { + private static final String SPACE_ID = "executor-lifecycle"; + private static final ByteString KEY = ByteString.copyFromUtf8("kept-key"); + private static final ByteString VALUE = ByteString.copyFromUtf8("kept-value"); + + @DataProvider + public Object[][] scenarios() { + return new Object[][] { + {false, "close"}, {false, "destroy"}, {false, "stop"}, + {true, "close"}, {true, "destroy"}, {true, "stop"}, + {false, "compactor"}, {true, "compactor"}, {true, "flusher"}, + {false, "constructor-failure"}, {false, "open-failure"}, {true, "open-failure"} + }; + } + + @Test(dataProvider = "scenarios") + public void executorLifecycle(boolean wal, String scenario) throws Exception { + Path root = Files.createTempDirectory("rocksdb-executor-"); + Path output = root.resolve("child.log"); + Process child = new ProcessBuilder(Path.of(System.getProperty("java.home"), "bin", "java").toString(), + "-ea", "-cp", System.getProperty("java.class.path"), RocksDBExecutorLifecycleTest.class.getName(), + Boolean.toString(wal), scenario, root.toString()) + .redirectErrorStream(true).redirectOutput(output.toFile()).start(); + try { + assertTrue(child.waitFor(45, TimeUnit.SECONDS), "Child did not exit: " + Files.readString(output)); + assertEquals(child.exitValue(), 0, Files.readString(output)); + } finally { + if (child.isAlive()) { + child.destroyForcibly(); + assertTrue(child.waitFor(10, TimeUnit.SECONDS)); + } + deleteDir(root); + } + } + + public static void main(String[] args) { + try { + boolean wal = Boolean.parseBoolean(args[0]); + String scenario = args[1]; + Path root = Path.of(args[2]); + if (scenario.endsWith("failure")) { + failedCreation(wal, scenario, root); + } else if (scenario.equals("compactor") || scenario.equals("flusher")) { + closeWaitsForTask(wal, scenario, root); + } else { + lifecycle(wal, scenario, root); + } + } catch (Throwable e) { + e.printStackTrace(); + System.exit(1); + } + // A successful child must exit without System.exit, even after compaction starts a worker. + } + + private static void lifecycle(boolean wal, String action, Path root) throws Exception { + SimpleMeterRegistry registry = new SimpleMeterRegistry(); + try { + Metrics.addRegistry(registry); + RocksDBKVEngine engine = newEngine(wal, root); + engine.start(); + IKVSpace space = engine.createIfMissing(SPACE_ID); + IKVSpace other = engine.createIfMissing("other"); + List names = wal ? List.of("compactor", "flusher") : List.of("compactor"); + List oldTimers = names.stream().map(name -> registry.get("kvspace.executor") + .tags("spaceId", SPACE_ID, "name", name).timer()).toList(); + List oldGauges = names.stream().map(name -> registry.get("kvspace.executor.pool.size") + .tags("spaceId", SPACE_ID, "name", name).gauge()).toList(); + exercise(space, registry); + exercise(other, registry); + List workers = workers(SPACE_ID); + assertEquals(workers.size(), wal ? 2 : 1); + assertEquals(executorMeters(registry, SPACE_ID), wal ? 18L : 9L); + if (action.equals("stop")) { + engine.stop(); + } else { + if (action.equals("destroy")) { + space.destroy(); + } else { + space.close(); + } + assertEquals(executorMeters(registry, "other"), wal ? 18L : 9L); + try (IKVSpaceReader reader = other.reader()) { + assertEquals(reader.get(KEY).orElseThrow(), VALUE); + } + } + assertEquals(executorMeters(registry, SPACE_ID), 0L); + assertStopped(workers); + if (!action.equals("stop")) { + engine.stop(); + } + assertEquals(executorMeters(registry, "other"), 0L); + + engine = newEngine(wal, root); + engine.start(); + IKVSpace replacement = engine.createIfMissing(SPACE_ID); + for (int i = 0; i < names.size(); i++) { + Timer newTimer = registry.get("kvspace.executor").tags("spaceId", SPACE_ID, "name", names.get(i)).timer(); + Gauge newGauge = registry.get("kvspace.executor.pool.size") + .tags("spaceId", SPACE_ID, "name", names.get(i)).gauge(); + assertNotSame(newTimer, oldTimers.get(i)); + assertNotSame(newGauge, oldGauges.get(i)); + assertEquals(newTimer.count(), 0L); + assertEquals(newGauge.value(), 0.0); + } + try (IKVSpaceReader reader = replacement.reader()) { + assertEquals(reader.get(KEY).isPresent(), !action.equals("destroy")); + if (!action.equals("destroy")) { + assertEquals(reader.get(KEY).orElseThrow(), VALUE); + } + } + exercise(replacement, registry); + for (String name : names) { + assertEquals(registry.get("kvspace.executor.pool.size") + .tags("spaceId", SPACE_ID, "name", name).gauge().value(), 1.0); + } + List replacementWorkers = workers(SPACE_ID); + engine.stop(); + assertStopped(replacementWorkers); + assertEquals(executorMeters(registry, SPACE_ID), 0L); + } finally { + Metrics.removeRegistry(registry); + registry.close(); + } + } + + private static void closeWaitsForTask(boolean wal, String worker, Path root) throws Exception { + BlockingRegistry registry = new BlockingRegistry(worker); + try { + Metrics.addRegistry(registry); + RocksDBKVEngine engine = newEngine(wal, root); + engine.start(); + IKVSpace space = engine.createIfMissing(SPACE_ID); + triggerCompaction(space); + if (worker.equals("flusher")) { + ((IWALableKVSpace) space).flush().get(10, TimeUnit.SECONDS); + } + assertTrue(registry.entered.await(10, TimeUnit.SECONDS)); + Path spaceDir = root.resolve("data").resolve(SPACE_ID); + Path dbDir = wal ? spaceDir : spaceDir.resolve( + Files.readString(spaceDir.resolve(RocksDBCPableKVSpace.ACTIVE_GEN_POINTER)).trim()); + CompletableFuture closed = CompletableFuture.runAsync(space::close); + try { + assertThrows(TimeoutException.class, () -> closed.get(300, TimeUnit.MILLISECONDS)); + try (Options options = new Options()) { + RocksDBException locked = org.testng.Assert.expectThrows(RocksDBException.class, + () -> RocksDB.open(options, dbDir.toString())); + assertTrue(locked.getMessage().toLowerCase().contains("lock"), locked.getMessage()); + } + } finally { + registry.release.countDown(); + } + closed.get(10, TimeUnit.SECONDS); + try (Options options = new Options(); RocksDB db = RocksDB.open(options, dbDir.toString())) { + assertTrue(db.isOwningHandle()); + } + engine.stop(); + } finally { + Metrics.removeRegistry(registry); + registry.close(); + } + } + + private static void failedCreation(boolean wal, String scenario, Path root) throws Exception { + SimpleMeterRegistry registry = new SimpleMeterRegistry(); + try { + Metrics.addRegistry(registry); + RocksDBKVEngine engine = newEngine(wal, root); + engine.start(); + if (scenario.equals("constructor-failure")) { + Path blocked = root.resolve("checkpoints").resolve(SPACE_ID); + Files.createFile(blocked); + assertThrows(Exception.class, () -> engine.createIfMissing(SPACE_ID)); + assertEquals(executorMeters(registry, SPACE_ID), 0L); + Files.delete(blocked); + } else { + Path spaceDir = root.resolve("data").resolve(SPACE_ID); + Path dbDir = wal ? spaceDir : spaceDir.resolve("generation"); + Files.createDirectories(dbDir); + if (!wal) { + Files.writeString(spaceDir.resolve(RocksDBCPableKVSpace.ACTIVE_GEN_POINTER), "generation"); + } + try (Options options = new Options().setCreateIfMissing(true); + RocksDB locked = RocksDB.open(options, dbDir.toString())) { + assertThrows(KVEngineException.class, () -> engine.createIfMissing(SPACE_ID)); + assertEquals(executorMeters(registry, SPACE_ID), 0L); + } + } + IKVSpace retried = engine.createIfMissing(SPACE_ID); + exercise(retried, registry); + engine.stop(); + assertEquals(executorMeters(registry, SPACE_ID), 0L); + } finally { + Metrics.removeRegistry(registry); + registry.close(); + } + } + + private static RocksDBKVEngine newEngine(boolean wal, Path root) { + Struct.Builder config = (wal ? RocksDBDefaultConfigs.WAL : RocksDBDefaultConfigs.CP).toBuilder() + .putFields(DB_ROOT_DIR, toValue(root.resolve("data").toString())) + .putFields(COMPACT_MIN_TOMBSTONE_KEYS, toValue(1)); + if (!wal) { + config.putFields(DB_CHECKPOINT_ROOT_DIR, toValue(root.resolve("checkpoints").toString())); + } + return wal ? new RocksDBWALableKVEngine(null, config.build()) + : new RocksDBCPableKVEngine(null, config.build()); + } + + private static void triggerCompaction(IKVSpace space) { + IKVSpaceWriter writer = space instanceof ICPableKVSpace cp ? cp.toWriter() + : ((IWALableKVSpace) space).toWriter(); + writer.put(KEY, VALUE) + .delete(ByteString.copyFromUtf8("deleted-1")) + .delete(ByteString.copyFromUtf8("deleted-2")).done(); + } + + private static void exercise(IKVSpace space, SimpleMeterRegistry registry) throws Exception { + triggerCompaction(space); + await().atMost(Duration.ofSeconds(10)).until(() -> registry.get("kvspace.executor") + .tags("spaceId", space.id(), "name", "compactor").timer().count() > 0); + if (space instanceof IWALableKVSpace wal) { + wal.flush().get(10, TimeUnit.SECONDS); + await().atMost(Duration.ofSeconds(10)).until(() -> registry.get("kvspace.executor") + .tags("spaceId", space.id(), "name", "flusher").timer().count() > 0); + } + } + + private static long executorMeters(SimpleMeterRegistry registry, String spaceId) { + return registry.getMeters().stream().map(Meter::getId) + .filter(id -> id.getName().startsWith("kvspace.executor") && spaceId.equals(id.getTag("spaceId"))).count(); + } + + private static List workers(String spaceId) { + return Thread.getAllStackTraces().keySet().stream().filter(thread -> thread.isAlive() + && (thread.getName().startsWith("kvspace-compactor-" + spaceId) + || thread.getName().startsWith("kvspace-flusher-" + spaceId))).toList(); + } + + private static void assertStopped(List workers) throws InterruptedException { + for (Thread worker : workers) { + worker.join(10000); + assertFalse(worker.isAlive(), worker.getName()); + } + } + + private static class BlockingRegistry extends SimpleMeterRegistry { + private final String worker; + private final CountDownLatch entered = new CountDownLatch(1); + private final CountDownLatch release = new CountDownLatch(1); + + BlockingRegistry(String worker) { + this.worker = worker; + } + + @Override + protected Timer newTimer(Meter.Id id, DistributionStatisticConfig config, PauseDetector pauseDetector) { + Timer timer = super.newTimer(id, config, pauseDetector); + if (id.getName().equals("kvspace.executor") && worker.equals(id.getTag("name"))) { + Timer monitored = spy(timer); + doAnswer(call -> { + entered.countDown(); + assertTrue(release.await(15, TimeUnit.SECONDS)); + return call.callRealMethod(); + }).when(monitored).record(anyLong(), any(TimeUnit.class)); + return monitored; + } + return timer; + } + } +} diff --git a/base-kv/base-kv-raft/src/main/java/org/apache/bifromq/basekv/raft/RaftNode.java b/base-kv/base-kv-raft/src/main/java/org/apache/bifromq/basekv/raft/RaftNode.java index 39d0e53b8..48881b611 100644 --- a/base-kv/base-kv-raft/src/main/java/org/apache/bifromq/basekv/raft/RaftNode.java +++ b/base-kv/base-kv-raft/src/main/java/org/apache/bifromq/basekv/raft/RaftNode.java @@ -24,7 +24,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.util.EnumMap; import java.util.HashMap; import java.util.List; @@ -40,6 +39,7 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import lombok.NonNull; +import org.apache.bifromq.base.util.MonitoredExecutors; import org.apache.bifromq.basekv.raft.event.RaftEvent; import org.apache.bifromq.basekv.raft.event.RaftEventType; import org.apache.bifromq.basekv.raft.exception.InternalError; @@ -73,7 +73,7 @@ public RaftNode(RaftConfig config, this.stateStorage = new MetricMonitoredStateStore(stateStore, Tags.of(tags)); this.id = stateStorage.local(); this.config = config.toBuilder().build(); - this.raftExecutor = ExecutorServiceMetrics.monitor(Metrics.globalRegistry, + this.raftExecutor = MonitoredExecutors.monitor(Metrics.globalRegistry, new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedTransferQueue<>(), threadFactory), id, "raft", Tags.of(tags)); stateStore.addStableListener(this::onStabilized); diff --git a/base-kv/base-kv-raft/src/test/java/org/apache/bifromq/basekv/raft/RaftNodeTest.java b/base-kv/base-kv-raft/src/test/java/org/apache/bifromq/basekv/raft/RaftNodeTest.java index 61a3d67ad..7ba5fc548 100644 --- a/base-kv/base-kv-raft/src/test/java/org/apache/bifromq/basekv/raft/RaftNodeTest.java +++ b/base-kv/base-kv-raft/src/test/java/org/apache/bifromq/basekv/raft/RaftNodeTest.java @@ -24,14 +24,16 @@ import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; -import org.apache.bifromq.basekv.raft.proto.ClusterConfig; -import org.apache.bifromq.basekv.raft.proto.Snapshot; import com.google.protobuf.ByteString; +import io.micrometer.core.instrument.Metrics; import java.util.Collections; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import java.util.function.Supplier; import lombok.extern.slf4j.Slf4j; +import org.apache.bifromq.basekv.raft.proto.ClusterConfig; +import org.apache.bifromq.basekv.raft.proto.Snapshot; import org.mockito.Mock; import org.mockito.MockitoAnnotations; import org.testng.annotations.AfterMethod; @@ -62,11 +64,31 @@ public void setup() { testNode = new RaftNode(new RaftConfig(), stateStorage, Executors.defaultThreadFactory()); } - @AfterMethod + @AfterMethod(alwaysRun = true) public void releaseMocks() throws Exception { + testNode.start(messageSender, eventListener, snapshotInstaller); + CompletableFuture stopped = testNode.stop(); + if (!stopped.isCancelled()) { + stopped.get(5, TimeUnit.SECONDS); + } + await().until(() -> executorMeters() == 0); closeable.close(); } + @Test + public void stopRemovesExecutorMeters() throws Exception { + assertEquals(executorMeters(), 9L); + testNode.start(messageSender, eventListener, snapshotInstaller); + testNode.stop().get(5, TimeUnit.SECONDS); + assertEquals(executorMeters(), 0L); + } + + private long executorMeters() { + return Metrics.globalRegistry.getMeters().stream() + .filter(meter -> meter.getId().getName().startsWith("raft.executor") + && "V1".equals(meter.getId().getTag("name"))).count(); + } + @Test public void id() { assertEquals(testNode.id(), "V1"); diff --git a/base-kv/base-kv-raft/src/test/java/org/apache/bifromq/basekv/raft/functest/RaftConfigMigrationTest.java b/base-kv/base-kv-raft/src/test/java/org/apache/bifromq/basekv/raft/functest/RaftConfigMigrationTest.java index ee8ec5d25..92012e916 100644 --- a/base-kv/base-kv-raft/src/test/java/org/apache/bifromq/basekv/raft/functest/RaftConfigMigrationTest.java +++ b/base-kv/base-kv-raft/src/test/java/org/apache/bifromq/basekv/raft/functest/RaftConfigMigrationTest.java @@ -45,29 +45,33 @@ public void testLeaderElectionWithPreVoteMigration() { put("V3", raftConfigWithoutPreVote); }}; RaftNodeGroup group = new RaftNodeGroup(clusterConfig(), raftConfigs); - // enough ticks for leader election - group.run(10, TimeUnit.MILLISECONDS); - await().until(() -> group.currentLeader().isPresent()); - String leader = group.currentLeader().get(); - assertTrue(group.awaitIndexCommitted(leader, 1)); - log.info("Leader {} elected", leader); + try { + // Start the ticker and wait for leader election. + group.run(10, TimeUnit.MILLISECONDS); + await().until(() -> group.currentLeader().isPresent()); + String leader = group.currentLeader().get(); + assertTrue(group.awaitIndexCommitted(leader, 1)); + log.info("Leader {} elected", leader); - String blockedFollower = group.currentFollowers().get(0); - log.info("Follower {} isolating", blockedFollower); - group.isolate(blockedFollower); - group.propose(leader, ByteString.copyFromUtf8("appCommand1")); - group.propose(leader, ByteString.copyFromUtf8("appCommand2")); - assertTrue(group.awaitIndexCommitted(leader, 3)); - assertTrue(group.awaitIndexCommitted(group.currentFollowers().get(1), 3)); + String blockedFollower = group.currentFollowers().get(0); + log.info("Follower {} isolating", blockedFollower); + group.isolate(blockedFollower); + group.propose(leader, ByteString.copyFromUtf8("appCommand1")); + group.propose(leader, ByteString.copyFromUtf8("appCommand2")); + assertTrue(group.awaitIndexCommitted(leader, 3)); + assertTrue(group.awaitIndexCommitted(group.currentFollowers().get(1), 3)); - // isolate leader, V3 and another follower will campaign, but V3 will not elected since that it has less logs - group.recoverNetwork(); - log.info("Unblock {} and isolating Leader {}", blockedFollower, leader); - group.isolate(leader); - await().until(() -> group.currentLeader().isPresent() && !leader.equals(group.currentLeader().get())); - String newLeader = group.currentLeader().get(); - log.info("New leader {} elected", newLeader); - assertNotEquals(newLeader, leader); - assertNotEquals(newLeader, blockedFollower); + // Restore the network, then isolate the leader to trigger another election. + group.recoverNetwork(); + log.info("Unblock {} and isolating Leader {}", blockedFollower, leader); + group.isolate(leader); + await().until(() -> group.currentLeader().isPresent() && !leader.equals(group.currentLeader().get())); + String newLeader = group.currentLeader().get(); + log.info("New leader {} elected", newLeader); + assertNotEquals(newLeader, leader); + assertNotEquals(newLeader, blockedFollower); + } finally { + group.shutdown(); + } } } diff --git a/base-kv/base-kv-store-server/src/main/java/org/apache/bifromq/basekv/store/range/KVRangeDumpSession.java b/base-kv/base-kv-store-server/src/main/java/org/apache/bifromq/basekv/store/range/KVRangeDumpSession.java index 96560a448..dd1ff4982 100644 --- a/base-kv/base-kv-store-server/src/main/java/org/apache/bifromq/basekv/store/range/KVRangeDumpSession.java +++ b/base-kv/base-kv-store-server/src/main/java/org/apache/bifromq/basekv/store/range/KVRangeDumpSession.java @@ -100,7 +100,7 @@ class KVRangeDumpSession { new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedTransferQueue<>(), EnvProvider.INSTANCE.newThreadFactory("basekv-snapshot-dumper")), - "mutator", "basekv.range", Tags.of(tags)); + "snapshot-dumper", "basekv.range", Tags.of(tags)); this.runner = new AsyncRunner("basekv.runner.sessiondump", executor); this.maxIdleDuration = maxIdleDuration; this.recorder = recorder; diff --git a/base-kv/base-kv-store-server/src/main/java/org/apache/bifromq/basekv/store/range/KVRangeFSM.java b/base-kv/base-kv-store-server/src/main/java/org/apache/bifromq/basekv/store/range/KVRangeFSM.java index c57d75ae4..b2959dc23 100644 --- a/base-kv/base-kv-store-server/src/main/java/org/apache/bifromq/basekv/store/range/KVRangeFSM.java +++ b/base-kv/base-kv-store-server/src/main/java/org/apache/bifromq/basekv/store/range/KVRangeFSM.java @@ -58,7 +58,6 @@ import com.google.protobuf.ByteString; import io.micrometer.core.instrument.Metrics; import io.micrometer.core.instrument.Tags; -import io.micrometer.core.instrument.binder.jvm.ExecutorServiceMetrics; import io.reactivex.rxjava3.core.Observable; import io.reactivex.rxjava3.core.ObservableSource; import io.reactivex.rxjava3.disposables.CompositeDisposable; @@ -92,6 +91,7 @@ import java.util.stream.Collectors; import org.apache.bifromq.base.util.AsyncRetry; import org.apache.bifromq.base.util.AsyncRunner; +import org.apache.bifromq.base.util.MonitoredExecutors; import org.apache.bifromq.baseenv.EnvProvider; import org.apache.bifromq.baseenv.ZeroCopyParser; import org.apache.bifromq.basehlc.HLC; @@ -241,12 +241,12 @@ public KVRangeFSM(String clusterId, this.metricManager = new KVRangeMetricManager(clusterId, hostStoreId, id); this.wal = new KVRangeWAL(clusterId, hostStoreId, id, walStore, opts.getWalRaftConfig(), opts.getMaxWALFatchBatchSize()); - this.fsmExecutor = ExecutorServiceMetrics.monitor(Metrics.globalRegistry, + this.fsmExecutor = MonitoredExecutors.monitor(Metrics.globalRegistry, new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedTransferQueue<>(), EnvProvider.INSTANCE.newThreadFactory("basekv-range-mutator")), "mutator", "basekv.range", Tags.of(tags)); - this.mgmtExecutor = ExecutorServiceMetrics.monitor(Metrics.globalRegistry, + this.mgmtExecutor = MonitoredExecutors.monitor(Metrics.globalRegistry, new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedTransferQueue<>(), EnvProvider.INSTANCE.newThreadFactory("basekv-range-manager")), @@ -497,13 +497,13 @@ public CompletableFuture destroy() { } return CompletableFuture.completedFuture(null); }, mgmtExecutor) + .thenCompose(v -> awaitShutdown(mgmtExecutor)) .whenComplete((v, e) -> { if (lifecycle.compareAndSet(Destroying, Destroyed)) { log.info("Range destroyed"); destroyedSignal.complete(null); } - }) - .thenCompose(v -> awaitShutdown(mgmtExecutor)); + }); } return destroyedSignal; } diff --git a/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/KVRangeStoreTest.java b/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/KVRangeStoreTest.java index 0bcbfd8fa..e6582fbdf 100644 --- a/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/KVRangeStoreTest.java +++ b/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/KVRangeStoreTest.java @@ -39,6 +39,7 @@ import com.google.common.collect.Sets; import com.google.protobuf.ByteString; import com.google.protobuf.Struct; +import io.micrometer.core.instrument.Metrics; import io.reactivex.rxjava3.core.Observable; import io.reactivex.rxjava3.subjects.PublishSubject; import java.io.File; @@ -178,6 +179,22 @@ protected void doTearDown(Method method) { log.info("Shutdown read task executor"); } + @Test + public void rangeExecutorMetersFollowSplitAndClose() { + KVRangeDescriptor first = firstRangeDescriptor(); + assertEquals(rangeExecutorMeters(), 18L); + rangeStore.split(first.getVer(), first.getId(), copyFromUtf8("split-key")).toCompletableFuture().join(); + await().atMost(Duration.ofSeconds(5)).until(() -> rangeExecutorMeters() == 36L); + rangeStore.stop(); + assertEquals(rangeExecutorMeters(), 0L); + } + + private long rangeExecutorMeters() { + return Metrics.globalRegistry.getMeters().stream() + .filter(meter -> meter.getId().getName().startsWith("basekv.range.executor") + && rangeStore.id().equals(meter.getId().getTag("storeId"))).count(); + } + @Test(groups = "integration") public void testKeyExist() { KVRangeDescriptor rangeDescriptor = firstRangeDescriptor(); diff --git a/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/range/KVRangeDumpSessionTest.java b/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/range/KVRangeDumpSessionTest.java index 6f3998782..c1cab83f0 100644 --- a/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/range/KVRangeDumpSessionTest.java +++ b/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/range/KVRangeDumpSessionTest.java @@ -32,11 +32,22 @@ import static org.testng.Assert.assertTrue; import com.google.protobuf.ByteString; +import io.micrometer.core.instrument.Metrics; +import io.micrometer.core.instrument.Tags; +import io.micrometer.core.instrument.Timer; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; import io.reactivex.rxjava3.subjects.PublishSubject; import java.time.Duration; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.LinkedTransferQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; +import org.apache.bifromq.base.util.MonitoredExecutors; import org.apache.bifromq.basekv.MockableTest; import org.apache.bifromq.basekv.proto.KVRangeId; import org.apache.bifromq.basekv.proto.KVRangeMessage; @@ -65,6 +76,73 @@ public class KVRangeDumpSessionTest extends MockableTest { @Mock private KVRangeDumpSession.DumpBytesRecorder dumpBytesRecorder; + @Test + public void parallelDumpsKeepTheirMetricsSeparateFromRangeExecutor() throws Exception { + SimpleMeterRegistry registry = new SimpleMeterRegistry(); + Metrics.addRegistry(registry); + KVRangeId rangeId = KVRangeIdUtil.generate(); + String[] tags = {"clusterId", "dump-test", "storeId", "store", "rangeId", KVRangeIdUtil.toString(rangeId)}; + ExecutorService rangeExecutor = MonitoredExecutors.monitor(Metrics.globalRegistry, + new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, new LinkedTransferQueue<>()), + "mutator", "basekv.range", Tags.of(tags)); + KVRangeSnapshot snapshot = KVRangeSnapshot.newBuilder().setId(rangeId).setCheckpointId("checkpoint").build(); + when(rangeAccessor.hasCheckpoint(snapshot)).thenReturn(true); + when(rangeAccessor.open(snapshot)).thenReturn(checkpointReader); + when(checkpointReader.iterator()).thenReturn(checkpointItr); + when(checkpointItr.isValid()).thenReturn(false); + when(messenger.receive()).thenReturn(PublishSubject.create()); + Set workers = ConcurrentHashMap.newKeySet(); + doAnswer(invocation -> { + KVRangeMessage message = invocation.getArgument(0); + if (message.getSaveSnapshotDataRequest().getFlag() == SaveSnapshotDataRequest.Flag.End) { + workers.add(Thread.currentThread()); + } + return null; + }).when(messenger).send(any()); + KVRangeDumpSession first = new KVRangeDumpSession("first", snapshot, rangeId, "peer", rangeAccessor, + messenger, Duration.ofSeconds(5), 1024, new SnapshotBandwidthGovernor(0), dumpBytesRecorder, tags); + KVRangeDumpSession second = new KVRangeDumpSession("second", snapshot, rangeId, "peer", rangeAccessor, + messenger, Duration.ofSeconds(5), 1024, new SnapshotBandwidthGovernor(0), dumpBytesRecorder, tags); + try { + verify(messenger, timeout(5000).times(2)).send(any()); + assertEquals(workers.size(), 2); + Timer dumpTimer = registry.get("basekv.range.executor").tags(tags) + .tag("name", "snapshot-dumper").timer(); + await().until(() -> dumpTimer.count() > 0); + assertEquals(registry.getMeters().stream() + .filter(meter -> "snapshot-dumper".equals(meter.getId().getTag("name")) + && KVRangeIdUtil.toString(rangeId).equals(meter.getId().getTag("rangeId"))).count(), 9L); + first.cancel(); + first.awaitDone().get(5, TimeUnit.SECONDS); + assertEquals(registry.get("basekv.range.executor").tags(tags) + .tag("name", "snapshot-dumper").timer(), dumpTimer); + final long dumpCount = dumpTimer.count(); + second.cancel(); + second.awaitDone().get(5, TimeUnit.SECONDS); + for (Thread worker : workers) { + worker.join(5000); + assertFalse(worker.isAlive()); + } + assertTrue(dumpTimer.count() > dumpCount); + Timer rangeTimer = registry.get("basekv.range.executor").tags(tags).tag("name", "mutator").timer(); + assertEquals(rangeTimer.count(), 0L); + rangeExecutor.submit(() -> 1).get(5, TimeUnit.SECONDS); + assertEquals(rangeTimer.count(), 1L); + } finally { + first.cancel(); + second.cancel(); + first.awaitDone().get(5, TimeUnit.SECONDS); + second.awaitDone().get(5, TimeUnit.SECONDS); + rangeExecutor.shutdownNow(); + assertTrue(rangeExecutor.awaitTermination(5, TimeUnit.SECONDS)); + Metrics.globalRegistry.getMeters().stream() + .filter(meter -> KVRangeIdUtil.toString(rangeId).equals(meter.getId().getTag("rangeId"))) + .toList().forEach(Metrics.globalRegistry::remove); + Metrics.removeRegistry(registry); + registry.close(); + } + } + @Test public void dumpEmptySnapshot() { String peerStoreId = "follower"; @@ -296,7 +374,7 @@ public void idle() { Thread.sleep(20); dumpSession.tick(); verify(messenger, timeout(100).times(2)).send(any()); - assertTrue(dumpSession.awaitDone().toCompletableFuture().isDone()); + assertEquals(dumpSession.awaitDone().get(5, TimeUnit.SECONDS), KVRangeDumpSession.Result.Canceled); await().untilAsserted(() -> { verify(checkpointItr).close(); verify(checkpointReader).close(); diff --git a/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/range/KVRangeFSMExecutorLifecycleTest.java b/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/range/KVRangeFSMExecutorLifecycleTest.java new file mode 100644 index 000000000..e22034d3a --- /dev/null +++ b/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/range/KVRangeFSMExecutorLifecycleTest.java @@ -0,0 +1,120 @@ +/* + * 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.basekv.store.range; + +import static org.apache.bifromq.basekv.localengine.StructUtil.toValue; +import static org.apache.bifromq.basekv.localengine.rocksdb.RocksDBDefaultConfigs.DB_ROOT_DIR; +import static org.apache.bifromq.basekv.utils.BoundaryUtil.FULL_BOUNDARY; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +import io.micrometer.core.instrument.Metrics; +import io.reactivex.rxjava3.core.Observable; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import org.apache.bifromq.basekv.TestCoProcFactory; +import org.apache.bifromq.basekv.localengine.rocksdb.RocksDBDefaultConfigs; +import org.apache.bifromq.basekv.proto.KVRangeId; +import org.apache.bifromq.basekv.proto.KVRangeSnapshot; +import org.apache.bifromq.basekv.proto.State; +import org.apache.bifromq.basekv.raft.proto.ClusterConfig; +import org.apache.bifromq.basekv.raft.proto.Snapshot; +import org.apache.bifromq.basekv.store.option.KVRangeOptions; +import org.apache.bifromq.basekv.store.wal.KVRangeWALStorageEngine; +import org.apache.bifromq.basekv.utils.KVRangeIdUtil; +import org.testng.annotations.Test; + +/** + * Checks that range destruction releases executor meters. + */ +public class KVRangeFSMExecutorLifecycleTest extends AbstractKVRangeTest { + @Test + public void destroyReleasesMetersBeforeCompletion() throws Exception { + KVRangeWALStorageEngine walEngine = new KVRangeWALStorageEngine("fsm-test", null, "rocksdb", + RocksDBDefaultConfigs.WAL.toBuilder() + .putFields(DB_ROOT_DIR, toValue(dbRootDir.resolve("wal").toString())).build()); + walEngine.start(); + ExecutorService auxiliary = Executors.newSingleThreadExecutor(); + CompletableFuture destroying = new CompletableFuture<>(); + CompletableFuture release = new CompletableFuture<>(); + KVRangeId id = KVRangeIdUtil.generate(); + ClusterConfig config = ClusterConfig.newBuilder().addVoters(walEngine.id()).build(); + KVRangeSnapshot snapshot = KVRangeSnapshot.newBuilder().setId(id).setVer(0).setLastAppliedIndex(0) + .setBoundary(FULL_BOUNDARY).setState(State.newBuilder().setType(State.StateType.Normal)) + .setClusterConfig(config).build(); + IKVRange range = spy(new KVRange(id, kvEngine.createIfMissing(KVRangeIdUtil.toString(id)), snapshot)); + doAnswer(invocation -> { + destroying.complete(null); + release.join(); + return invocation.callRealMethod(); + }).when(range).destroy(); + KVRangeFSM fsm = fsm(range, snapshot, walEngine, auxiliary); + IKVRangeMessenger messenger = mock(IKVRangeMessenger.class); + when(messenger.receive()).thenReturn(Observable.never()); + try { + fsm.open(messenger).get(5, TimeUnit.SECONDS); + assertEquals(executorMeters(id), 18L); + CompletableFuture destroyed = fsm.destroy(); + destroying.get(5, TimeUnit.SECONDS); + CompletableFuture metersAtCompletion = destroyed.thenApply(v -> executorMeters(id)); + release.complete(null); + assertEquals(metersAtCompletion.get(5, TimeUnit.SECONDS), Long.valueOf(0)); + IKVRange replacementRange = new KVRange(id, + kvEngine.createIfMissing(KVRangeIdUtil.toString(id)), snapshot); + KVRangeFSM replacement = fsm(replacementRange, snapshot, walEngine, auxiliary); + try { + replacement.open(messenger).get(5, TimeUnit.SECONDS); + assertEquals(executorMeters(id), 18L); + } finally { + replacement.close().get(5, TimeUnit.SECONDS); + } + assertEquals(executorMeters(id), 0L); + } finally { + release.complete(null); + fsm.close().get(5, TimeUnit.SECONDS); + walEngine.stop(); + auxiliary.shutdown(); + assertTrue(auxiliary.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private KVRangeFSM fsm(IKVRange range, KVRangeSnapshot snapshot, KVRangeWALStorageEngine walEngine, + ExecutorService auxiliary) { + Snapshot walSnapshot = Snapshot.newBuilder().setClusterConfig(snapshot.getClusterConfig()) + .setData(snapshot.toByteString()).build(); + return new KVRangeFSM("fsm-test", walEngine.id(), snapshot.getId(), new TestCoProcFactory(), range, + walEngine.create(snapshot.getId(), walSnapshot), auxiliary, auxiliary, new KVRangeOptions(), + List.of(), (closed, reset) -> { }, "clusterId", "fsm-test", "storeId", walEngine.id(), + "rangeId", KVRangeIdUtil.toString(snapshot.getId())); + } + + private long executorMeters(KVRangeId id) { + return Metrics.globalRegistry.getMeters().stream() + .filter(meter -> meter.getId().getName().startsWith("basekv.range.executor") + && KVRangeIdUtil.toString(id).equals(meter.getId().getTag("rangeId"))).count(); + } +} diff --git a/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/wal/KVRangeWALTest.java b/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/wal/KVRangeWALTest.java index 526229a40..6ef541f59 100644 --- a/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/wal/KVRangeWALTest.java +++ b/base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/wal/KVRangeWALTest.java @@ -92,7 +92,12 @@ public void doTearDown(Method method) { public void testId() { KVRangeWAL wal = new KVRangeWAL("testcluster", replicaId, id, raftStateStorage, config, 1024); - assertEquals(wal.storeId(), replicaId); + try { + assertEquals(wal.storeId(), replicaId); + } finally { + wal.start(); + wal.close().join(); + } } @SneakyThrows diff --git a/base-util/pom.xml b/base-util/pom.xml index 8eaeb9142..af4d01f3d 100644 --- a/base-util/pom.xml +++ b/base-util/pom.xml @@ -44,6 +44,11 @@ slf4j-api + + io.micrometer + micrometer-registry-prometheus + test + org.awaitility awaitility diff --git a/base-util/src/main/java/org/apache/bifromq/base/util/MonitoredExecutors.java b/base-util/src/main/java/org/apache/bifromq/base/util/MonitoredExecutors.java new file mode 100644 index 000000000..9c02bfa2f --- /dev/null +++ b/base-util/src/main/java/org/apache/bifromq/base/util/MonitoredExecutors.java @@ -0,0 +1,431 @@ +/* + * 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.base.util; + +import static java.util.concurrent.TimeUnit.NANOSECONDS; + +import io.micrometer.core.instrument.Clock; +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.FunctionCounter; +import io.micrometer.core.instrument.Gauge; +import io.micrometer.core.instrument.Meter; +import io.micrometer.core.instrument.MeterRegistry; +import io.micrometer.core.instrument.Tag; +import io.micrometer.core.instrument.Tags; +import io.micrometer.core.instrument.Timer; +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.Future; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +/** + * Registers executor meters and removes them when the executor shuts down. + * + *

The caller must own the executor and use the returned service for tasks and shutdown. + * Shutdown removes meters before queued tasks finish. Use awaitTermination to wait for those tasks. + * Each registry, prefix, and tag set must have one owner. External registrations and meter filters + * must not merge the meter identities of different owners. + * + *

Meter definitions follow Micrometer 1.11.12 ExecutorServiceMetrics (Copyright 2017 VMware, Inc.), + * under the Apache License, Version 2.0. Task timing follows its TimedExecutorService and + * TimedScheduledExecutorService (Copyright 2017, 2019 VMware, Inc.). + * This implementation adds meter ownership and shutdown cleanup. + * See https://github.com/micrometer-metrics/micrometer/tree/v1.11.12/micrometer-core. + */ +public final class MonitoredExecutors { + // Registration and removal share this lock. Task submission and execution do not use it. + private static final Map> OWNERS = new IdentityHashMap<>(); + + private MonitoredExecutors() { + } + + /** + * Monitors an executor without a prefix or caller tags. + */ + public static ExecutorService monitor(MeterRegistry registry, ExecutorService executor, String name) { + return monitor(registry, executor, name, "", Tags.empty()); + } + + /** + * Monitors an executor with caller tags and no prefix. + */ + public static ExecutorService monitor(MeterRegistry registry, ExecutorService executor, String name, + Iterable tags) { + return monitor(registry, executor, name, "", tags); + } + + /** + * Monitors an executor with a prefix and no caller tags. + */ + public static ExecutorService monitor(MeterRegistry registry, ExecutorService executor, String name, + String metricPrefix) { + return monitor(registry, executor, name, metricPrefix, Tags.empty()); + } + + /** + * Monitors a ThreadPoolExecutor, ScheduledThreadPoolExecutor, or owned ForkJoinPool. + * + *

Unsupported executors and duplicate identities fail before registration. Registration failure + * rolls back the acquired meters. The caller remains responsible for the supplied executor. + */ + public static ExecutorService monitor(MeterRegistry registry, ExecutorService executor, String name, + String metricPrefix, Iterable tags) { + if (!(executor instanceof ThreadPoolExecutor) && !(executor instanceof ForkJoinPool)) { + throw new IllegalArgumentException("Executor type is not supported: " + executor.getClass().getName()); + } + String prefix = metricPrefix == null || metricPrefix.isBlank() ? "" + : metricPrefix.endsWith(".") ? metricPrefix : metricPrefix + "."; + Identity identity = new Identity(prefix, Tags.concat(tags, "name", name)); + synchronized (OWNERS) { + Map registryOwners = OWNERS.computeIfAbsent(registry, k -> new HashMap<>()); + if (registryOwners.containsKey(identity)) { + throw new IllegalArgumentException("Executor meters already have an owner: " + identity); + } + Binding binding = new Binding(registry, identity); + registryOwners.put(identity, binding); + try { + binding.register(executor); + Timing timing = new Timing(registry.config().clock(), binding.executionTimer, binding.idleTimer); + if (executor instanceof ScheduledThreadPoolExecutor scheduled) { + return new MonitoredScheduler(scheduled, binding, timing); + } + return new MonitoredExecutor(executor, binding, timing); + } catch (RuntimeException failure) { + binding.removeMeters(); + throw failure; + } + } + } + + /** + * Monitors a scheduled executor without a prefix or caller tags. + */ + public static ScheduledExecutorService monitor(MeterRegistry registry, ScheduledExecutorService executor, + String name) { + return monitor(registry, executor, name, "", Tags.empty()); + } + + /** + * Monitors a scheduled executor with caller tags and no prefix. + */ + public static ScheduledExecutorService monitor(MeterRegistry registry, ScheduledExecutorService executor, + String name, Iterable tags) { + return monitor(registry, executor, name, "", tags); + } + + /** + * Monitors a scheduled executor with a prefix and no caller tags. + */ + public static ScheduledExecutorService monitor(MeterRegistry registry, ScheduledExecutorService executor, + String name, String metricPrefix) { + return monitor(registry, executor, name, metricPrefix, Tags.empty()); + } + + /** + * Monitors a scheduled executor and preserves its scheduling interface. + */ + public static ScheduledExecutorService monitor(MeterRegistry registry, ScheduledExecutorService executor, + String name, String metricPrefix, Iterable tags) { + return (ScheduledExecutorService) monitor(registry, (ExecutorService) executor, name, metricPrefix, tags); + } + + private record Identity(String prefix, Tags tags) { + } + + private static final class Binding { + private final MeterRegistry registry; + private final Identity identity; + private final List meters = new ArrayList<>(11); + private Timer executionTimer; + private Timer idleTimer; + private Counter scheduledOnce; + private Counter scheduledRepetitively; + private boolean removed; + + private Binding(MeterRegistry registry, Identity identity) { + this.registry = registry; + this.identity = identity; + } + + private void register(ExecutorService executor) { + String prefix = identity.prefix(); + Tags tags = identity.tags(); + if (executor instanceof ThreadPoolExecutor pool) { + meters.add(FunctionCounter.builder(prefix + "executor.completed", pool, + ThreadPoolExecutor::getCompletedTaskCount) + .tags(tags) + .description("The approximate total number of tasks that have completed execution") + .baseUnit("tasks").register(registry)); + meters.add(Gauge.builder(prefix + "executor.active", pool, ThreadPoolExecutor::getActiveCount) + .tags(tags) + .description("The approximate number of threads that are actively executing tasks") + .baseUnit("threads").register(registry)); + meters.add(Gauge.builder(prefix + "executor.queued", pool, p -> p.getQueue().size()) + .tags(tags) + .description("The approximate number of tasks that are queued for execution") + .baseUnit("tasks").register(registry)); + meters.add(Gauge.builder(prefix + "executor.queue.remaining", pool, + p -> p.getQueue().remainingCapacity()) + .tags(tags) + .description("The number of additional elements that this queue can ideally accept without " + + "blocking") + .baseUnit("tasks").register(registry)); + meters.add(Gauge.builder(prefix + "executor.pool.size", pool, ThreadPoolExecutor::getPoolSize) + .tags(tags).description("The current number of threads in the pool") + .baseUnit("threads").register(registry)); + meters.add(Gauge.builder(prefix + "executor.pool.core", pool, ThreadPoolExecutor::getCorePoolSize) + .tags(tags).description("The core number of threads for the pool") + .baseUnit("threads").register(registry)); + meters.add(Gauge.builder(prefix + "executor.pool.max", pool, ThreadPoolExecutor::getMaximumPoolSize) + .tags(tags).description("The maximum allowed number of threads in the pool") + .baseUnit("threads").register(registry)); + } else { + ForkJoinPool pool = (ForkJoinPool) executor; + meters.add(FunctionCounter.builder(prefix + "executor.steals", pool, ForkJoinPool::getStealCount) + .tags(tags) + .description("Estimate of the total number of tasks stolen from " + + "one thread's work queue by another. The reported value " + + "underestimates the actual total number of steals when the pool is not quiescent") + .register(registry)); + meters.add(Gauge.builder(prefix + "executor.queued", pool, ForkJoinPool::getQueuedTaskCount) + .tags(tags) + .description("An estimate of the total number of tasks currently held in queues by worker threads") + .register(registry)); + meters.add(Gauge.builder(prefix + "executor.active", pool, ForkJoinPool::getActiveThreadCount) + .tags(tags) + .description("An estimate of the number of threads that are currently stealing or executing tasks") + .register(registry)); + meters.add(Gauge.builder(prefix + "executor.running", pool, ForkJoinPool::getRunningThreadCount) + .tags(tags) + .description("An estimate of the number of worker threads that are not blocked waiting to join " + + "tasks or for other managed synchronization threads") + .register(registry)); + } + executionTimer = Timer.builder(prefix + "executor").tags(tags).register(registry); + meters.add(executionTimer); + idleTimer = Timer.builder(prefix + "executor.idle").tags(tags).register(registry); + meters.add(idleTimer); + if (executor instanceof ScheduledThreadPoolExecutor) { + scheduledOnce = Counter.builder(prefix + "executor.scheduled.once").tags(tags).register(registry); + meters.add(scheduledOnce); + scheduledRepetitively = Counter.builder(prefix + "executor.scheduled.repetitively") + .tags(tags).register(registry); + meters.add(scheduledRepetitively); + } + } + + private void close() { + synchronized (OWNERS) { + removeMeters(); + } + } + + // The caller holds OWNERS until removal finishes and the identity becomes available. + private void removeMeters() { + if (removed) { + return; + } + meters.forEach(registry::remove); + meters.clear(); + Map registryOwners = OWNERS.get(registry); + registryOwners.remove(identity); + if (registryOwners.isEmpty()) { + OWNERS.remove(registry); + } + removed = true; + } + } + + private static class MonitoredExecutor implements ExecutorService { + private final ExecutorService delegate; + private final Binding binding; + private final Timing timing; + + private MonitoredExecutor(ExecutorService delegate, Binding binding, Timing timing) { + this.delegate = delegate; + this.binding = binding; + this.timing = timing; + } + + @Override + public void shutdown() { + delegate.shutdown(); + binding.close(); + } + + @Override + public List shutdownNow() { + List pending = delegate.shutdownNow(); + binding.close(); + return pending; + } + + @Override + public boolean isShutdown() { + return delegate.isShutdown(); + } + + @Override + public boolean isTerminated() { + return delegate.isTerminated(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + return delegate.awaitTermination(timeout, unit); + } + + @Override + public void execute(Runnable command) { + delegate.execute(timing.wrap(command)); + } + + @Override + public Future submit(Runnable task) { + return delegate.submit(timing.wrap(task)); + } + + @Override + public Future submit(Runnable task, T result) { + return delegate.submit(timing.wrap(task), result); + } + + @Override + public Future submit(Callable task) { + return delegate.submit(timing.wrap(task)); + } + + @Override + public List> invokeAll(Collection> tasks) throws InterruptedException { + return delegate.invokeAll(wrapAll(tasks)); + } + + @Override + public List> invokeAll(Collection> tasks, long timeout, TimeUnit unit) + throws InterruptedException { + return delegate.invokeAll(wrapAll(tasks), timeout, unit); + } + + @Override + public T invokeAny(Collection> tasks) + throws InterruptedException, ExecutionException { + return delegate.invokeAny(wrapAll(tasks)); + } + + @Override + public T invokeAny(Collection> tasks, long timeout, TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + return delegate.invokeAny(wrapAll(tasks), timeout, unit); + } + + private Collection> wrapAll(Collection> tasks) { + List> wrapped = new ArrayList<>(tasks.size()); + for (Callable task : tasks) { + wrapped.add(timing.wrap(task)); + } + return wrapped; + } + } + + private static final class MonitoredScheduler extends MonitoredExecutor implements ScheduledExecutorService { + private final ScheduledExecutorService scheduler; + private final Timer executionTimer; + private final Counter scheduledOnce; + private final Counter scheduledRepetitively; + + private MonitoredScheduler(ScheduledExecutorService scheduler, Binding binding, Timing timing) { + super(scheduler, binding, timing); + this.scheduler = scheduler; + executionTimer = binding.executionTimer; + scheduledOnce = binding.scheduledOnce; + scheduledRepetitively = binding.scheduledRepetitively; + } + + @Override + public ScheduledFuture schedule(Runnable command, long delay, TimeUnit unit) { + scheduledOnce.increment(); + return scheduler.schedule(executionTimer.wrap(command), delay, unit); + } + + @Override + public ScheduledFuture schedule(Callable callable, long delay, TimeUnit unit) { + scheduledOnce.increment(); + return scheduler.schedule(executionTimer.wrap(callable), delay, unit); + } + + @Override + public ScheduledFuture scheduleAtFixedRate(Runnable command, long initialDelay, long period, + TimeUnit unit) { + scheduledRepetitively.increment(); + return scheduler.scheduleAtFixedRate(executionTimer.wrap(command), initialDelay, period, unit); + } + + @Override + public ScheduledFuture scheduleWithFixedDelay(Runnable command, long initialDelay, long delay, + TimeUnit unit) { + scheduledRepetitively.increment(); + return scheduler.scheduleWithFixedDelay(executionTimer.wrap(command), initialDelay, delay, unit); + } + } + + private record Timing(Clock clock, Timer executionTimer, Timer idleTimer) { + private Runnable wrap(Runnable task) { + Objects.requireNonNull(task, "task"); + long submitted = clock.monotonicTime(); + return () -> { + idleTimer.record(clock.monotonicTime() - submitted, NANOSECONDS); + long started = clock.monotonicTime(); + try { + task.run(); + } finally { + executionTimer.record(clock.monotonicTime() - started, NANOSECONDS); + } + }; + } + + private Callable wrap(Callable task) { + Objects.requireNonNull(task, "task"); + long submitted = clock.monotonicTime(); + return () -> { + idleTimer.record(clock.monotonicTime() - submitted, NANOSECONDS); + long started = clock.monotonicTime(); + try { + return task.call(); + } finally { + executionTimer.record(clock.monotonicTime() - started, NANOSECONDS); + } + }; + } + } +} diff --git a/base-util/src/test/java/org/apache/bifromq/base/util/MonitoredExecutorsExportTest.java b/base-util/src/test/java/org/apache/bifromq/base/util/MonitoredExecutorsExportTest.java new file mode 100644 index 000000000..61c5d18ba --- /dev/null +++ b/base-util/src/test/java/org/apache/bifromq/base/util/MonitoredExecutorsExportTest.java @@ -0,0 +1,128 @@ +/* + * 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.base.util; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + +import io.micrometer.core.instrument.Meter; +import io.micrometer.core.instrument.Tags; +import io.micrometer.core.instrument.composite.CompositeMeterRegistry; +import io.micrometer.core.instrument.config.MeterFilter; +import io.micrometer.core.instrument.distribution.DistributionStatisticConfig; +import io.micrometer.prometheus.PrometheusConfig; +import io.micrometer.prometheus.PrometheusMeterRegistry; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +/** + * Checks executor meter removal from Prometheus exports. + */ +public class MonitoredExecutorsExportTest { + @DataProvider + public Object[][] poolKinds() { + return new Object[][] {{0, 9}, {1, 11}, {2, 6}}; + } + + @Test(dataProvider = "poolKinds") + public void removesExportedSamples(int kind, int count) throws Exception { + PrometheusMeterRegistry registry = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT); + try { + ExecutorService service = MonitoredExecutors.monitor(registry, pool(kind), "first", "owned", Tags.empty()); + ExecutorService other = MonitoredExecutors.monitor(registry, pool(kind), "other", "owned", Tags.empty()); + try { + service.submit(() -> 1).get(5, TimeUnit.SECONDS); + other.submit(() -> 2).get(5, TimeUnit.SECONDS); + assertEquals(registry.getMeters().size(), count * 2); + assertTrue(registry.scrape().contains("name=\"first\"")); + assertTrue(registry.scrape().contains("name=\"other\"")); + service.shutdown(); + assertTrue(service.awaitTermination(5, TimeUnit.SECONDS)); + String scrape = registry.scrape(); + assertFalse(scrape.contains("name=\"first\"")); + assertTrue(scrape.contains("name=\"other\"")); + assertEquals(other.submit(() -> 3).get(5, TimeUnit.SECONDS), Integer.valueOf(3)); + } finally { + service.shutdownNow(); + other.shutdownNow(); + assertTrue(service.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(other.awaitTermination(5, TimeUnit.SECONDS)); + } + assertFalse(registry.scrape().contains("owned_executor")); + } finally { + registry.close(); + } + } + + @Test + public void followsCompositeMembershipAndKeepsHistograms() throws Exception { + CompositeMeterRegistry registry = new CompositeMeterRegistry(); + PrometheusMeterRegistry first = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT); + PrometheusMeterRegistry second = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT); + try { + second.config().meterFilter(new MeterFilter() { + @Override + public DistributionStatisticConfig configure(Meter.Id id, DistributionStatisticConfig config) { + return DistributionStatisticConfig.builder().percentilesHistogram(true).build().merge(config); + } + }); + ExecutorService service = MonitoredExecutors.monitor(registry, pool(0), "composite"); + try { + registry.add(first); + service.submit(() -> 1).get(5, TimeUnit.SECONDS); + assertTrue(first.scrape().contains("executor_seconds_count")); + registry.add(second); + service.submit(() -> 2).get(5, TimeUnit.SECONDS); + assertTrue(second.scrape().contains("executor_seconds_bucket")); + registry.remove(first); + long firstCount = first.get("executor").timer().count(); + service.submit(() -> 3).get(5, TimeUnit.SECONDS); + assertEquals(first.get("executor").timer().count(), firstCount); + assertEquals(second.get("executor").timer().count(), 2L); + service.shutdown(); + assertTrue(service.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(registry.getMeters().isEmpty()); + assertFalse(second.scrape().contains("name=\"composite\"")); + } finally { + service.shutdownNow(); + assertTrue(service.awaitTermination(5, TimeUnit.SECONDS)); + } + } finally { + registry.close(); + first.close(); + second.close(); + } + } + + private static ExecutorService pool(int kind) { + return switch (kind) { + case 0 -> new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, new LinkedBlockingQueue<>()); + case 1 -> new ScheduledThreadPoolExecutor(1); + default -> new ForkJoinPool(1); + }; + } +} diff --git a/base-util/src/test/java/org/apache/bifromq/base/util/MonitoredExecutorsTest.java b/base-util/src/test/java/org/apache/bifromq/base/util/MonitoredExecutorsTest.java new file mode 100644 index 000000000..3ac61d8e6 --- /dev/null +++ b/base-util/src/test/java/org/apache/bifromq/base/util/MonitoredExecutorsTest.java @@ -0,0 +1,652 @@ +/* + * 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.base.util; + +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.Gauge; +import io.micrometer.core.instrument.Meter; +import io.micrometer.core.instrument.MockClock; +import io.micrometer.core.instrument.Tags; +import io.micrometer.core.instrument.Timer; +import io.micrometer.core.instrument.binder.jvm.ExecutorServiceMetrics; +import io.micrometer.core.instrument.config.MeterFilter; +import io.micrometer.core.instrument.simple.SimpleConfig; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.TreeMap; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +public class MonitoredExecutorsTest { + private final List executors = new ArrayList<>(); + private final List registries = new ArrayList<>(); + private SimpleMeterRegistry registry; + + @BeforeMethod + public void setup() { + registry = registry(new SimpleMeterRegistry()); + } + + @AfterMethod(alwaysRun = true) + public void cleanup() throws InterruptedException { + for (int i = executors.size() - 1; i >= 0; i--) { + executors.get(i).shutdownNow(); + } + for (ExecutorService executor : executors) { + assertTrue(executor.awaitTermination(5, SECONDS)); + } + registries.forEach(SimpleMeterRegistry::close); + executors.clear(); + registries.clear(); + } + + @DataProvider + public Object[][] poolKinds() { + return new Object[][] {{0, 9}, {1, 11}, {2, 6}}; + } + + @DataProvider + public Object[][] shutdownPaths() { + return new Object[][] {{0, 9, false}, {0, 9, true}, {1, 11, false}, {1, 11, true}, + {2, 6, false}, {2, 6, true}}; + } + + @Test(dataProvider = "shutdownPaths") + public void removesAllMeters(int kind, int expected, boolean immediate) throws Exception { + ExecutorService service = monitor(pool(kind), "test"); + service.submit(() -> 7).get(5, SECONDS); + assertEquals(registry.getMeters().size(), expected); + if (immediate) { + service.shutdownNow(); + } else { + service.shutdown(); + } + assertTrue(service.isShutdown()); + assertTrue(service.awaitTermination(5, SECONDS)); + assertTrue(service.isTerminated()); + assertTrue(registry.getMeters().isEmpty()); + service.shutdown(); + service.shutdownNow(); + assertTrue(registry.getMeters().isEmpty()); + } + + @Test(dataProvider = "poolKinds") + public void keepsMicrometerMeterDefinitions(int kind, int count) { + ExecutorService raw = pool(kind); + SimpleMeterRegistry baseline = registry(new SimpleMeterRegistry()); + ExecutorService legacy = watch(ExecutorServiceMetrics.monitor(baseline, raw, "test", "prefix", + Tags.of("store", "s", "name", "ignored"))); + ExecutorService current = watch(MonitoredExecutors.monitor(registry, raw, "test", "prefix.", + Tags.of("name", "ignored", "store", "s"))); + assertEquals(meterDefinitions(registry), meterDefinitions(baseline)); + assertEquals(registry.getMeters().size(), count); + assertEquals(current instanceof ScheduledExecutorService, kind == 1); + legacy.shutdown(); + current.shutdown(); + } + + @Test + public void normalizesPrefixAndTagsBeforeClaimingIdentity() { + ExecutorService first = watch(MonitoredExecutors.monitor(registry, pool(0), "test", "prefix", + Tags.of("b", "2", "a", "1"))); + ExecutorService candidate = pool(1); + expectThrows(IllegalArgumentException.class, () -> MonitoredExecutors.monitor(registry, candidate, "test", + "prefix.", Tags.of("a", "1", "b", "2"))); + assertFalse(candidate.isShutdown()); + assertEquals(registry.getMeters().size(), 9); + first.shutdown(); + watch(MonitoredExecutors.monitor(registry, candidate, "test", "prefix.", Tags.of("b", "2", "a", "1"))); + assertEquals(registry.getMeters().size(), 11); + } + + @Test + public void normalizesEmptyPrefixes() { + ExecutorService first = watch(MonitoredExecutors.monitor(registry, pool(0), "test", null, Tags.empty())); + ExecutorService candidate = pool(0); + expectThrows(IllegalArgumentException.class, () -> MonitoredExecutors.monitor(registry, candidate, "test", + " ", Tags.empty())); + first.shutdown(); + watch(MonitoredExecutors.monitor(registry, candidate, "test", Tags.empty())); + assertNotNull(registry.find("executor").timer()); + } + + @Test + public void isolatesRegistriesNamesTagsAndUnrelatedMeters() throws Exception { + SimpleMeterRegistry otherRegistry = registry(new SimpleMeterRegistry()); + ExecutorService first = monitor(pool(0), "first"); + ExecutorService second = monitor(pool(0), "second"); + ExecutorService other = watch(MonitoredExecutors.monitor(otherRegistry, pool(0), "first")); + ExecutorService tagged = watch(MonitoredExecutors.monitor(registry, pool(0), "first", Tags.of("store", "s"))); + Counter business = registry.counter("executor.business", "name", "first"); + first.shutdown(); + second.submit(() -> 1).get(5, SECONDS); + other.submit(() -> 1).get(5, SECONDS); + tagged.submit(() -> 1).get(5, SECONDS); + business.increment(); + assertEquals(registry.get("executor").tag("name", "second").timer().count(), 1L); + assertEquals(registry.get("executor").tag("store", "s").timer().count(), 1L); + assertEquals(otherRegistry.get("executor").timer().count(), 1L); + assertEquals(business.count(), 1.0); + assertEquals(registry.getMeters().size(), 19); + } + + @Test + public void rejectsUnsupportedAndRepeatedWrappersBeforeRegistration() { + ExecutorService decorated = watch(Executors.newSingleThreadExecutor()); + expectThrows(IllegalArgumentException.class, () -> MonitoredExecutors.monitor(registry, decorated, "decorated")); + assertTrue(registry.getMeters().isEmpty()); + assertFalse(decorated.isShutdown()); + ExecutorService first = monitor(pool(0), "first"); + expectThrows(IllegalArgumentException.class, () -> MonitoredExecutors.monitor(registry, first, "second")); + assertEquals(registry.getMeters().size(), 9); + } + + @Test + public void oldTasksAndRepeatedShutdownCannotAffectReplacement() throws Exception { + ExecutorService old = monitor(pool(0), "test"); + Timer oldTimer = registry.get("executor").timer(); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch finish = new CountDownLatch(1); + Future running = old.submit(() -> { + started.countDown(); + await(finish); + }); + await(started); + old.shutdown(); + assertFalse(old.isTerminated()); + assertFalse(old.awaitTermination(0, SECONDS)); + assertTrue(registry.getMeters().isEmpty()); + ThreadPoolExecutor rawReplacement = watch(new ThreadPoolExecutor(3, 3, 0, SECONDS, + new LinkedBlockingQueue<>())); + ExecutorService replacement = monitor(rawReplacement, "test"); + Timer newTimer = registry.get("executor").timer(); + assertEquals(registry.get("executor.pool.core").gauge().value(), 3.0); + assertEquals(newTimer.count(), 0L); + finish.countDown(); + running.get(5, SECONDS); + assertTrue(old.awaitTermination(5, SECONDS)); + assertEquals(oldTimer.count(), 1L); + assertEquals(newTimer.count(), 0L); + old.shutdownNow(); + old.shutdown(); + assertSame(registry.get("executor").timer(), newTimer); + replacement.submit(() -> 1).get(5, SECONDS); + assertEquals(newTimer.count(), 1L); + } + + @Test + public void shutdownNowStillInterruptsAfterShutdown() throws Exception { + ExecutorService service = monitor(pool(0), "test"); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + service.submit(() -> { + started.countDown(); + try { + new CountDownLatch(1).await(); + } catch (InterruptedException e) { + interrupted.countDown(); + Thread.currentThread().interrupt(); + } + }); + await(started); + service.shutdown(); + service.shutdownNow(); + await(interrupted); + assertTrue(service.awaitTermination(5, SECONDS)); + } + + @Test + public void rejectsConcurrentOwners() throws Exception { + ExecutorService callers = watch(Executors.newFixedThreadPool(2)); + ExecutorService firstPool = pool(0); + ExecutorService secondPool = pool(0); + CyclicBarrier start = new CyclicBarrier(2); + List> results = new ArrayList<>(); + for (ExecutorService raw : List.of(firstPool, secondPool)) { + results.add(callers.submit(() -> { + start.await(5, SECONDS); + try { + return MonitoredExecutors.monitor(registry, raw, "same"); + } catch (IllegalArgumentException e) { + return null; + } + })); + } + int owners = 0; + for (Future result : results) { + ExecutorService service = result.get(5, SECONDS); + if (service != null) { + watch(service); + owners++; + } + } + assertEquals(owners, 1); + assertEquals(registry.getMeters().size(), 9); + } + + @Test + public void doesNotReleaseIdentityBeforeRemovalCompletes() throws Exception { + ExecutorService old = monitor(pool(0), "test"); + ExecutorService replacementPool = pool(0); + ExecutorService callers = watch(Executors.newFixedThreadPool(2)); + CountDownLatch removing = new CountDownLatch(1); + CountDownLatch allowRemoval = new CountDownLatch(1); + AtomicInteger removed = new AtomicInteger(); + registry.config().onMeterRemoved(meter -> { + if (removed.getAndIncrement() == 0) { + removing.countDown(); + await(allowRemoval); + } + }); + Future closing = callers.submit(old::shutdown); + await(removing); + CountDownLatch registering = new CountDownLatch(1); + Future creating = callers.submit(() -> { + registering.countDown(); + return MonitoredExecutors.monitor(registry, replacementPool, "test"); + }); + try { + await(registering); + assertFalse(creating.isDone()); + } finally { + allowRemoval.countDown(); + } + closing.get(5, SECONDS); + ExecutorService replacement = watch(creating.get(5, SECONDS)); + assertEquals(registry.getMeters().size(), 9); + replacement.submit(() -> 1).get(5, SECONDS); + old.shutdownNow(); + assertEquals(registry.get("executor").timer().count(), 1L); + } + + @Test + public void closesDifferentOwnersConcurrently() throws Exception { + ExecutorService callers = watch(Executors.newFixedThreadPool(3)); + CyclicBarrier registered = new CyclicBarrier(3); + List> results = new ArrayList<>(); + for (int kind = 0; kind < 3; kind++) { + ExecutorService raw = pool(kind); + String name = "owner-" + kind; + results.add(callers.submit(() -> { + ExecutorService service = MonitoredExecutors.monitor(registry, raw, name); + try { + registered.await(5, SECONDS); + assertEquals(service.submit(() -> name).get(5, SECONDS), name); + } finally { + service.shutdown(); + } + return null; + })); + } + for (Future result : results) { + result.get(5, SECONDS); + } + assertTrue(registry.getMeters().isEmpty()); + } + + @Test(dataProvider = "poolKinds") + public void rollsBackRegistrationWithoutClosingCallerPool(int kind, int expectedMeters) throws Exception { + Gauge conflict = Gauge.builder("executor.idle", () -> 1).tag("name", "test").register(registry); + ExecutorService raw = pool(kind); + assertEquals(raw.submit(() -> 1).get(5, SECONDS), Integer.valueOf(1)); + expectThrows(IllegalArgumentException.class, () -> MonitoredExecutors.monitor(registry, raw, "test")); + assertEquals(registry.getMeters(), List.of(conflict)); + assertFalse(raw.isShutdown()); + assertEquals(raw.submit(() -> 2).get(5, SECONDS), Integer.valueOf(2)); + registry.remove(conflict); + watch(MonitoredExecutors.monitor(registry, raw, "test")); + assertEquals(registry.getMeters().size(), expectedMeters); + } + + @Test + public void preservesPendingList() throws Exception { + PendingExecutor raw = watch(new PendingExecutor()); + ExecutorService service = monitor(raw, "test"); + CountDownLatch started = new CountDownLatch(1); + service.submit(() -> { + started.countDown(); + new CountDownLatch(1).await(); + return null; + }); + await(started); + Future pendingTask = service.submit(() -> 1); + List pending = service.shutdownNow(); + assertSame(pending, raw.pending); + assertEquals(pending.size(), 1); + assertSame(pending.get(0), pendingTask); + assertFalse(pendingTask.isCancelled()); + assertFalse(pendingTask.isDone()); + assertTrue(registry.getMeters().isEmpty()); + monitor(pool(0), "test"); + assertEquals(registry.getMeters().size(), 9); + } + + @Test + public void rejectsNullTasksBeforeSubmission() { + ExecutorService service = monitor(pool(0), "test"); + expectThrows(NullPointerException.class, () -> service.execute(null)); + expectThrows(NullPointerException.class, () -> service.submit((Runnable) null)); + expectThrows(NullPointerException.class, () -> service.submit((Callable) null)); + assertEquals(registry.get("executor").timer().count(), 0L); + assertEquals(registry.get("executor.pool.size").gauge().value(), 0.0); + } + + @Test + public void supportsDeniedMeters() { + registry.config().meterFilter(MeterFilter.deny()); + ExecutorService first = monitor(pool(0), "test"); + assertTrue(registry.getMeters().isEmpty()); + first.shutdown(); + monitor(pool(0), "test").shutdownNow(); + } + + @Test(dataProvider = "poolKinds") + public void forwardsSubmitResultsAndTaskFailures(int kind, int ignored) throws Exception { + ExecutorService service = monitor(pool(kind), "test"); + AtomicInteger runs = new AtomicInteger(); + assertEquals(service.submit((Callable) () -> 7).get(5, SECONDS), Integer.valueOf(7)); + assertEquals(service.submit(runs::incrementAndGet, "result").get(5, SECONDS), "result"); + assertEquals(service.submit((Runnable) runs::incrementAndGet).get(5, SECONDS), null); + IllegalArgumentException failure = new IllegalArgumentException("Task failed"); + ExecutionException thrown = expectThrows(ExecutionException.class, + () -> service.submit(() -> { throw failure; }).get(5, SECONDS)); + // ForkJoinPool can copy task exceptions when another thread reads the result. + assertTrue(thrown.getCause() instanceof IllegalArgumentException); + assertEquals(runs.get(), 2); + assertEquals(registry.get("executor").timer().count(), 4L); + assertEquals(registry.get("executor.idle").timer().count(), 4L); + } + + @Test + public void recordsWaitingAndExecutionSeparately() throws Exception { + MockClock clock = new MockClock(); + SimpleMeterRegistry timedRegistry = registry(new SimpleMeterRegistry(SimpleConfig.DEFAULT, clock)); + ThreadPoolExecutor raw = (ThreadPoolExecutor) pool(0); + CountDownLatch blocked = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + raw.execute(() -> { blocked.countDown(); await(release); }); + await(blocked); + ExecutorService service = watch(MonitoredExecutors.monitor(timedRegistry, raw, "test")); + Future task = service.submit(() -> clock.add(Duration.ofSeconds(2))); + clock.add(Duration.ofSeconds(3)); + release.countDown(); + task.get(5, SECONDS); + assertEquals(timedRegistry.get("executor.idle").timer().totalTime(SECONDS), 3.0); + assertEquals(timedRegistry.get("executor").timer().totalTime(SECONDS), 2.0); + } + + @Test + public void forwardsCancellationAndCallerRunsRejectionPolicy() throws Exception { + ThreadPoolExecutor raw = watch(new ThreadPoolExecutor(1, 1, 0, SECONDS, new ArrayBlockingQueue<>(1), + new ThreadPoolExecutor.CallerRunsPolicy())); + ExecutorService service = monitor(raw, "test"); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + Future running = service.submit(() -> { + started.countDown(); + try { + new CountDownLatch(1).await(); + } catch (InterruptedException e) { + interrupted.countDown(); + Thread.currentThread().interrupt(); + } + }); + await(started); + Future queued = service.submit(() -> 2); + Thread caller = Thread.currentThread(); + assertSame(service.submit(Thread::currentThread).get(5, SECONDS), caller); + assertEquals(registry.get("executor").timer().count(), 1L); + assertTrue(queued.cancel(false)); + assertTrue(running.cancel(true)); + await(interrupted); + assertTrue(running.isCancelled()); + } + + @Test + public void forwardsBatchResultsTimeoutsAndInterrupts() throws Exception { + ExecutorService service = monitor(pool(0), "test"); + List> tasks = List.of(() -> 1, () -> 2); + List> all = service.invokeAll(tasks); + assertEquals(all.get(0).get(), Integer.valueOf(1)); + assertEquals(all.get(1).get(), Integer.valueOf(2)); + assertEquals(service.invokeAll(tasks, 5, SECONDS).get(1).get(), Integer.valueOf(2)); + assertEquals(service.invokeAny(List.of(() -> { throw new IllegalStateException(); }, () -> 3)), Integer.valueOf(3)); + assertEquals(service.invokeAny(tasks, 5, SECONDS), Integer.valueOf(1)); + Callable blocked = () -> { new CountDownLatch(1).await(); return 0; }; + List> timedOut = service.invokeAll(List.of(blocked, blocked), 20, MILLISECONDS); + assertTrue(timedOut.stream().allMatch(Future::isCancelled)); + expectThrows(TimeoutException.class, () -> service.invokeAny(List.of(blocked), 20, MILLISECONDS)); + Thread.currentThread().interrupt(); + try { + expectThrows(InterruptedException.class, () -> service.invokeAll(List.of(blocked))); + } finally { + Thread.interrupted(); + } + service.shutdown(); + expectThrows(RejectedExecutionException.class, () -> service.submit(() -> 1)); + } + + @Test + public void schedulesWithoutCountingDelayAsIdleTime() throws Exception { + MockClock clock = new MockClock(); + SimpleMeterRegistry timedRegistry = registry(new SimpleMeterRegistry(SimpleConfig.DEFAULT, clock)); + ScheduledThreadPoolExecutor raw = watch(new ScheduledThreadPoolExecutor(1)); + CountDownLatch blocked = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + raw.execute(() -> { blocked.countDown(); await(release); }); + await(blocked); + ScheduledExecutorService service = watch(MonitoredExecutors.monitor(timedRegistry, raw, "test")); + ScheduledFuture task = service.schedule(() -> { + clock.add(Duration.ofSeconds(2)); + return 7; + }, 1, MILLISECONDS); + clock.add(Duration.ofHours(1)); + release.countDown(); + assertEquals(task.get(5, SECONDS), Integer.valueOf(7)); + service.schedule(() -> { }, 0, MILLISECONDS).get(5, SECONDS); + assertEquals(timedRegistry.get("executor.idle").timer().count(), 0L); + assertEquals(timedRegistry.get("executor").timer().totalTime(SECONDS), 2.0); + assertEquals(timedRegistry.get("executor.scheduled.once").counter().count(), 2.0); + } + + @DataProvider + public Object[][] periodicKinds() { + return new Object[][] {{false}, {true}}; + } + + @Test(dataProvider = "periodicKinds") + public void preservesPeriodicFailureAndCancellation(boolean fixedRate) throws Exception { + ScheduledExecutorService service = watch(MonitoredExecutors.monitor(registry, + watch(new ScheduledThreadPoolExecutor(1)), "test", Tags.empty())); + AtomicInteger calls = new AtomicInteger(); + Runnable failing = () -> { + if (calls.incrementAndGet() == 3) { + throw new IllegalStateException("Periodic task failed"); + } + }; + ScheduledFuture future = fixedRate ? service.scheduleAtFixedRate(failing, 0, 1, MILLISECONDS) + : service.scheduleWithFixedDelay(failing, 0, 1, MILLISECONDS); + expectThrows(ExecutionException.class, () -> future.get(5, SECONDS)); + assertEquals(calls.get(), 3); + assertEquals(registry.get("executor").timer().count(), 3L); + assertEquals(registry.get("executor.idle").timer().count(), 0L); + ScheduledFuture cancelled = fixedRate ? service.scheduleAtFixedRate(failing, 1, 1, SECONDS) + : service.scheduleWithFixedDelay(failing, 1, 1, SECONDS); + assertTrue(cancelled.cancel(false)); + assertTrue(cancelled.isCancelled()); + assertEquals(registry.get("executor.scheduled.repetitively").counter().count(), 2.0); + } + + @Test + public void countsRejectedSchedulingCallsWithoutRegisteringMetersAgain() { + ScheduledExecutorService service = watch(MonitoredExecutors.monitor(registry, + watch(new ScheduledThreadPoolExecutor(1)), "test")); + Counter once = registry.get("executor.scheduled.once").counter(); + Counter repeating = registry.get("executor.scheduled.repetitively").counter(); + service.shutdown(); + expectThrows(RejectedExecutionException.class, () -> service.schedule(() -> { }, 0, SECONDS)); + expectThrows(RejectedExecutionException.class, () -> service.schedule(() -> 1, 0, SECONDS)); + expectThrows(RejectedExecutionException.class, () -> service.scheduleAtFixedRate(() -> { }, 0, 1, SECONDS)); + expectThrows(RejectedExecutionException.class, () -> service.scheduleWithFixedDelay(() -> { }, 0, 1, SECONDS)); + assertEquals(once.count(), 2.0); + assertEquals(repeating.count(), 2.0); + assertTrue(registry.getMeters().isEmpty()); + } + + @Test + public void supportsPrefixOnlyOverloads() { + ExecutorService regular = watch(MonitoredExecutors.monitor(registry, pool(0), "test", "regular")); + ScheduledExecutorService scheduled = watch(MonitoredExecutors.monitor(registry, + watch(new ScheduledThreadPoolExecutor(1)), "test", "scheduled")); + assertNotNull(registry.find("regular.executor").timer()); + assertNotNull(registry.find("scheduled.executor.scheduled.once").counter()); + regular.shutdown(); + scheduled.shutdown(); + assertTrue(registry.getMeters().isEmpty()); + } + + @Test + public void executeKeepsUncaughtExceptionBehaviorAndRecordsTime() throws Exception { + CompletableFuture uncaught = new CompletableFuture<>(); + ThreadPoolExecutor raw = watch(new ThreadPoolExecutor(1, 1, 0, SECONDS, new LinkedBlockingQueue<>(), task -> { + Thread worker = new Thread(task); + worker.setUncaughtExceptionHandler((thread, error) -> uncaught.complete(error)); + return worker; + })); + ExecutorService service = monitor(raw, "test"); + IllegalStateException failure = new IllegalStateException("Task failed"); + service.execute(() -> { throw failure; }); + assertSame(uncaught.get(5, SECONDS), failure); + assertEquals(registry.get("executor").timer().count(), 1L); + service.submit(() -> 1).get(5, SECONDS); + assertEquals(registry.get("executor").timer().count(), 2L); + } + + @Test + public void batchFailuresKeepTheirResults() throws Exception { + ExecutorService service = monitor(pool(0), "test"); + Callable failing = () -> { throw new IllegalStateException("Task failed"); }; + List> results = service.invokeAll(List.of(failing, () -> 7)); + expectThrows(ExecutionException.class, () -> results.get(0).get()); + assertEquals(results.get(1).get(), Integer.valueOf(7)); + expectThrows(ExecutionException.class, () -> service.invokeAny(List.of(failing))); + expectThrows(ExecutionException.class, () -> service.invokeAny(List.of(failing), 5, SECONDS)); + } + + @Test + public void rollsBackTimersWhenScheduledCounterRegistrationFails() { + Gauge conflict = Gauge.builder("executor.scheduled.repetitively", () -> 1) + .tag("name", "test").register(registry); + ScheduledThreadPoolExecutor raw = watch(new ScheduledThreadPoolExecutor(1)); + expectThrows(IllegalArgumentException.class, () -> MonitoredExecutors.monitor(registry, raw, "test")); + assertEquals(registry.getMeters(), List.of(conflict)); + registry.remove(conflict); + watch(MonitoredExecutors.monitor(registry, raw, "test")); + assertEquals(registry.getMeters().size(), 11); + } + + private ExecutorService monitor(ExecutorService raw, String name) { + return watch(MonitoredExecutors.monitor(registry, raw, name)); + } + + private ExecutorService pool(int kind) { + return switch (kind) { + case 1 -> watch(new ScheduledThreadPoolExecutor(1)); + case 2 -> watch(new ForkJoinPool(1)); + default -> watch(new ThreadPoolExecutor(1, 1, 0, SECONDS, new LinkedBlockingQueue<>())); + }; + } + + private T watch(T executor) { + executors.add(executor); + return executor; + } + + private T registry(T meterRegistry) { + registries.add(meterRegistry); + return meterRegistry; + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, SECONDS), "Task did not reach the expected state"); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + + private static Map> meterDefinitions(SimpleMeterRegistry registry) { + Map> definitions = new TreeMap<>(); + for (Meter meter : registry.getMeters()) { + Meter.Id id = meter.getId(); + definitions.put(id.getName(), Arrays.asList(id.getType(), id.getTags(), id.getBaseUnit(), id.getDescription())); + } + return definitions; + } + + private static final class PendingExecutor extends ThreadPoolExecutor { + private List pending; + + private PendingExecutor() { + super(1, 1, 0, SECONDS, new LinkedBlockingQueue<>()); + } + + @Override + public List shutdownNow() { + pending = super.shutdownNow(); + return pending; + } + } + +} diff --git a/bifromq-dist/bifromq-dist-worker/pom.xml b/bifromq-dist/bifromq-dist-worker/pom.xml index bd77847cb..8f7c3ff6f 100644 --- a/bifromq-dist/bifromq-dist-worker/pom.xml +++ b/bifromq-dist/bifromq-dist-worker/pom.xml @@ -30,6 +30,10 @@ bifromq-dist-worker + + org.apache.bifromq + base-util + org.apache.bifromq base-env-provider diff --git a/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DeliverExecutor.java b/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DeliverExecutor.java index 0dacb110f..3f68681c1 100644 --- a/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DeliverExecutor.java +++ b/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DeliverExecutor.java @@ -22,7 +22,7 @@ import static org.apache.bifromq.plugin.eventcollector.ThreadLocalEventPool.getLocal; import io.micrometer.core.instrument.Metrics; -import io.micrometer.core.instrument.binder.jvm.ExecutorServiceMetrics; +import io.micrometer.core.instrument.Tags; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ExecutorService; import java.util.concurrent.LinkedTransferQueue; @@ -30,6 +30,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import lombok.extern.slf4j.Slf4j; +import org.apache.bifromq.base.util.MonitoredExecutors; import org.apache.bifromq.baseenv.EnvProvider; import org.apache.bifromq.deliverer.DeliveryCall; import org.apache.bifromq.deliverer.IMessageDeliverer; @@ -48,13 +49,14 @@ public class DeliverExecutor { private final ConcurrentLinkedQueue tasks = new ConcurrentLinkedQueue<>(); private final AtomicBoolean sending = new AtomicBoolean(); - public DeliverExecutor(int id, IMessageDeliverer deliverer, IEventCollector eventCollector) { + public DeliverExecutor(int id, IMessageDeliverer deliverer, IEventCollector eventCollector, String... tags) { this.eventCollector = eventCollector; this.deliverer = deliverer; - executor = ExecutorServiceMetrics.monitor(Metrics.globalRegistry, + executor = MonitoredExecutors.monitor(Metrics.globalRegistry, new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedTransferQueue<>(), - EnvProvider.INSTANCE.newThreadFactory("deliver-executor-" + id)), Integer.toString(id), "deliver"); + EnvProvider.INSTANCE.newThreadFactory("deliver-executor-" + id)), + Integer.toString(id), "deliver", Tags.of(tags)); } public void submit(NormalMatching route, TopicMessagePackHolder msgPackHolder, boolean inline) { diff --git a/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DeliverExecutorGroup.java b/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DeliverExecutorGroup.java index fb6a0ac69..6ef331795 100644 --- a/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DeliverExecutorGroup.java +++ b/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DeliverExecutorGroup.java @@ -76,7 +76,8 @@ class DeliverExecutorGroup implements IDeliverExecutorGroup { IResourceThrottler resourceThrottler, ISettingProvider settingProvider, int fanoutParallelism, - int inlineFanOutThreshold) { + int inlineFanOutThreshold, + String... tags) { this.eventCollector = eventCollector; this.resourceThrottler = resourceThrottler; this.settingProvider = settingProvider; @@ -96,7 +97,7 @@ class DeliverExecutorGroup implements IDeliverExecutorGroup { .build()); fanoutExecutors = new DeliverExecutor[fanoutParallelism]; for (int i = 0; i < fanoutParallelism; i++) { - fanoutExecutors[i] = new DeliverExecutor(i, deliverer, eventCollector); + fanoutExecutors[i] = new DeliverExecutor(i, deliverer, eventCollector, tags); } } diff --git a/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DistWorkerCoProcFactory.java b/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DistWorkerCoProcFactory.java index a24b3e4dd..4980d4df5 100644 --- a/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DistWorkerCoProcFactory.java +++ b/bifromq-dist/bifromq-dist-worker/src/main/java/org/apache/bifromq/dist/worker/DistWorkerCoProcFactory.java @@ -94,7 +94,8 @@ public IKVRangeCoProc createCoProc(String clusterId, String storeId, KVRangeId i "clusterId", clusterId, "storeId", storeId, "rangeId", KVRangeIdUtil.toString(id)); IDeliverExecutorGroup deliverExecutorGroup = new DeliverExecutorGroup( - deliverer, eventCollector, resourceThrottler, settingProvider, fanoutParallelism, inlineFanOutThreshold); + deliverer, eventCollector, resourceThrottler, settingProvider, fanoutParallelism, inlineFanOutThreshold, + "clusterId", clusterId, "storeId", storeId, "rangeId", KVRangeIdUtil.toString(id)); return new DistWorkerCoProc( id, rangeReaderProvider, routeCache, tenantsState, deliverExecutorGroup, subscriptionChecker); } diff --git a/bifromq-dist/bifromq-dist-worker/src/test/java/org/apache/bifromq/dist/worker/DeliverExecutorLifecycleTest.java b/bifromq-dist/bifromq-dist-worker/src/test/java/org/apache/bifromq/dist/worker/DeliverExecutorLifecycleTest.java new file mode 100644 index 000000000..51ebd0ec0 --- /dev/null +++ b/bifromq-dist/bifromq-dist-worker/src/test/java/org/apache/bifromq/dist/worker/DeliverExecutorLifecycleTest.java @@ -0,0 +1,87 @@ +/* + * 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.dist.worker; + +import static org.apache.bifromq.dist.worker.schema.cache.Matchings.normalMatching; +import static org.awaitility.Awaitility.await; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; + +import io.micrometer.core.instrument.Metrics; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import java.time.Duration; +import java.util.concurrent.CompletableFuture; +import org.apache.bifromq.deliverer.DeliveryCallResult; +import org.apache.bifromq.deliverer.IMessageDeliverer; +import org.apache.bifromq.deliverer.TopicMessagePackHolder; +import org.apache.bifromq.plugin.eventcollector.IEventCollector; +import org.apache.bifromq.type.TopicMessagePack; +import org.testng.annotations.Test; + +/** + * Checks executor meter ownership across ranges. + */ +public class DeliverExecutorLifecycleTest { + @Test + public void sameIndexInDifferentRangesHasSeparateMeters() { + SimpleMeterRegistry registry = new SimpleMeterRegistry(); + Metrics.addRegistry(registry); + IMessageDeliverer deliverer = mock(IMessageDeliverer.class); + when(deliverer.schedule(any())).thenReturn(CompletableFuture.completedFuture(DeliveryCallResult.OK)); + IEventCollector collector = mock(IEventCollector.class); + DeliverExecutor first = new DeliverExecutor(0, deliverer, collector, "rangeId", "first"); + DeliverExecutor second = new DeliverExecutor(0, deliverer, collector, "rangeId", "second"); + try { + assertEquals(count(registry, "first"), 9L); + assertEquals(count(registry, "second"), 9L); + submit(first); + submit(second); + await().atMost(Duration.ofSeconds(5)).until(() -> registry.get("deliver.executor") + .tags("rangeId", "first").timer().count() > 0); + await().atMost(Duration.ofSeconds(5)).until(() -> registry.get("deliver.executor") + .tags("rangeId", "second").timer().count() > 0); + first.shutdown(); + assertEquals(count(registry, "first"), 0L); + long before = registry.get("deliver.executor").tags("rangeId", "second").timer().count(); + submit(second); + await().atMost(Duration.ofSeconds(5)).until(() -> registry.get("deliver.executor") + .tags("rangeId", "second").timer().count() > before); + assertEquals(count(registry, "second"), 9L); + } finally { + first.shutdown(); + second.shutdown(); + assertEquals(count(registry, "first") + count(registry, "second"), 0L); + Metrics.removeRegistry(registry); + registry.close(); + } + } + + private static void submit(DeliverExecutor executor) { + executor.submit(normalMatching("tenant", "topic", 0, "receiver", "deliverer", 1), + TopicMessagePackHolder.hold(TopicMessagePack.newBuilder().setTopic("topic").build()), false); + } + + private static long count(SimpleMeterRegistry registry, String range) { + return registry.getMeters().stream().filter(meter -> meter.getId().getName().startsWith("deliver.executor") + && range.equals(meter.getId().getTag("rangeId"))).count(); + } +} diff --git a/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/InboxStoreCoProc.java b/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/InboxStoreCoProc.java index 4a2440ac7..460892850 100644 --- a/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/InboxStoreCoProc.java +++ b/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/InboxStoreCoProc.java @@ -202,7 +202,7 @@ final class InboxStoreCoProc implements IKVRangeCoProc { this.tenantStats = new TenantsStats(rangeReaderProvider, "clusterId", clusterId, "storeId", storeId, "rangeId", KVRangeIdUtil.toString(id)); this.delayTaskRunner = new DelayTaskRunner<>(TenantInboxInstance::compareTo, HLC.INST::getPhysical, - expireRateLimit); + expireRateLimit, "clusterId", clusterId, "storeId", storeId, "rangeId", KVRangeIdUtil.toString(id)); this.detachTimeout = detachTimeout; } diff --git a/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunner.java b/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunner.java index b6144cb98..14f21761d 100644 --- a/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunner.java +++ b/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunner.java @@ -19,23 +19,27 @@ package org.apache.bifromq.inbox.store.delay; -import org.apache.bifromq.baseenv.EnvProvider; +import static com.google.common.util.concurrent.Uninterruptibles.awaitTerminationUninterruptibly; + import com.google.common.annotations.VisibleForTesting; import com.google.common.util.concurrent.RateLimiter; import io.micrometer.core.instrument.Metrics; -import io.micrometer.core.instrument.binder.jvm.ExecutorServiceMetrics; +import io.micrometer.core.instrument.Tags; import java.time.Duration; import java.util.Comparator; import java.util.Map; import java.util.NavigableMap; import java.util.Set; import java.util.TreeMap; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.function.Supplier; +import org.apache.bifromq.base.util.MonitoredExecutors; +import org.apache.bifromq.baseenv.EnvProvider; /** * DelayTaskRunner is a utility class to schedule tasks to be executed after a specified delay. @@ -60,17 +64,19 @@ public class DelayTaskRunner> implements IDelayTas * @param comparator the comparator to be used for sorting the keys * @param currentMillisSupplier a supplier to get the current time in milliseconds * @param rateLimit the rate limit for triggering the scheduled tasks + * @param tags the range metric tags */ public DelayTaskRunner(Comparator comparator, Supplier currentMillisSupplier, - int rateLimit) { + int rateLimit, + String... tags) { this.currentMillisSupplier = currentMillisSupplier; this.sortedDeadlines = new TreeMap<>( Comparator.comparingLong((SortKey sk) -> sk.deadlineTS).thenComparing(sk -> sk.key, comparator)); this.rateLimiter = RateLimiter.create(rateLimit); - executor = ExecutorServiceMetrics.monitor(Metrics.globalRegistry, + executor = MonitoredExecutors.monitor(Metrics.globalRegistry, new ScheduledThreadPoolExecutor(1, EnvProvider.INSTANCE.newThreadFactory("delay-task-runner")), - "delay-task-runner"); + "delay-task-runner", "inbox.delay", Tags.of(tags)); } @VisibleForTesting @@ -184,10 +190,14 @@ public void cancelAll(Set keys) { @Override public void shutdown() { - executor.submit(() -> { + CompletableFuture.runAsync(() -> { isShutdown = true; + if (triggerTask != null) { + triggerTask.cancel(false); + } executor.shutdown(); - }); + }, executor).join(); + awaitTerminationUninterruptibly(executor); } private void trigger() { diff --git a/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/delay/IDelayTaskRunner.java b/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/delay/IDelayTaskRunner.java index dd3894ffc..adc07d599 100644 --- a/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/delay/IDelayTaskRunner.java +++ b/bifromq-inbox/bifromq-inbox-store/src/main/java/org/apache/bifromq/inbox/store/delay/IDelayTaskRunner.java @@ -53,7 +53,7 @@ public interface IDelayTaskRunner { void cancelAll(Set keys); /** - * Shutdown the DelayTaskRunner. + * Shut down the runner and wait for its executor to terminate. */ void shutdown(); } diff --git a/bifromq-inbox/bifromq-inbox-store/src/test/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunnerLifecycleTest.java b/bifromq-inbox/bifromq-inbox-store/src/test/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunnerLifecycleTest.java new file mode 100644 index 000000000..b9ba9af5e --- /dev/null +++ b/bifromq-inbox/bifromq-inbox-store/src/test/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunnerLifecycleTest.java @@ -0,0 +1,148 @@ +/* + * 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.inbox.store.delay; + +import static org.awaitility.Awaitility.await; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertThrows; + +import io.micrometer.core.instrument.Metrics; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import org.testng.annotations.Test; + +/** + * Checks executor meter removal before range replacement. + */ +public class DelayTaskRunnerLifecycleTest { + @Test + public void closesOnlyTheOwningRange() throws Exception { + SimpleMeterRegistry registry = new SimpleMeterRegistry(); + Metrics.addRegistry(registry); + List> runners = new ArrayList<>(); + CompletableFuture release = new CompletableFuture<>(); + try { + DelayTaskRunner first = runner("first"); + runners.add(first); + DelayTaskRunner second = runner("second"); + runners.add(second); + assertEquals(count(registry, "first"), 11L); + assertEquals(count(registry, "second"), 11L); + execute(first).get(5, TimeUnit.SECONDS); + execute(second).get(5, TimeUnit.SECONDS); + final Thread worker = execute(first, release).get(5, TimeUnit.SECONDS); + CompletableFuture closeStarted = new CompletableFuture<>(); + CompletableFuture closing = CompletableFuture.runAsync(() -> { + closeStarted.complete(null); + first.shutdown(); + }); + runners.remove(first); + closeStarted.get(5, TimeUnit.SECONDS); + assertThrows(TimeoutException.class, () -> closing.get(100, TimeUnit.MILLISECONDS)); + assertEquals(count(registry, "first"), 11L); + release.complete(null); + closing.get(5, TimeUnit.SECONDS); + worker.join(5000); + assertFalse(worker.isAlive()); + assertEquals(count(registry, "first"), 0L); + assertEquals(count(registry, "second"), 11L); + execute(second).get(5, TimeUnit.SECONDS); + DelayTaskRunner replacement = runner("first"); + runners.add(replacement); + execute(replacement).get(5, TimeUnit.SECONDS); + assertEquals(count(registry, "first"), 11L); + } finally { + release.complete(null); + runners.forEach(DelayTaskRunner::shutdown); + await().atMost(Duration.ofSeconds(5)) + .until(() -> count(registry, "first") + count(registry, "second") == 0); + Metrics.removeRegistry(registry); + registry.close(); + } + } + + @Test + public void shutdownCancelsFutureTriggerAndStopsWorker() throws Exception { + DelayTaskRunner runner = runner("future-trigger"); + CompletableFuture scheduled = new CompletableFuture<>(); + AtomicBoolean executed = new AtomicBoolean(); + Thread worker; + try { + runner.schedule("future", new IDelayedTask<>() { + @Override + public Duration getDelay() { + scheduled.complete(Thread.currentThread()); + return Duration.ofSeconds(1); + } + + @Override + public CompletableFuture run(String key, IDelayTaskRunner owner) { + executed.set(true); + return CompletableFuture.completedFuture(null); + } + }); + worker = scheduled.get(5, TimeUnit.SECONDS); + } finally { + runner.shutdown(); + } + worker.join(5000); + assertFalse(worker.isAlive()); + assertFalse(executed.get()); + } + + private static DelayTaskRunner runner(String range) { + return new DelayTaskRunner<>(String::compareTo, System::currentTimeMillis, 1000, + "clusterId", "delay-test", "storeId", "store", "rangeId", range); + } + + private static CompletableFuture execute(DelayTaskRunner runner) { + return execute(runner, CompletableFuture.completedFuture(null)); + } + + private static CompletableFuture execute(DelayTaskRunner runner, CompletableFuture release) { + CompletableFuture done = new CompletableFuture<>(); + runner.schedule("task", new IDelayedTask<>() { + @Override + public Duration getDelay() { + return Duration.ZERO; + } + + @Override + public CompletableFuture run(String key, IDelayTaskRunner owner) { + done.complete(Thread.currentThread()); + release.join(); + return CompletableFuture.completedFuture(null); + } + }); + return done; + } + + private static long count(SimpleMeterRegistry registry, String range) { + return registry.getMeters().stream().filter(meter -> meter.getId().getName().startsWith("inbox.delay.executor") + && range.equals(meter.getId().getTag("rangeId"))).count(); + } +} diff --git a/bifromq-inbox/bifromq-inbox-store/src/test/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunnerTest.java b/bifromq-inbox/bifromq-inbox-store/src/test/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunnerTest.java index 561908617..555405048 100644 --- a/bifromq-inbox/bifromq-inbox-store/src/test/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunnerTest.java +++ b/bifromq-inbox/bifromq-inbox-store/src/test/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunnerTest.java @@ -28,6 +28,7 @@ import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; +import io.micrometer.core.instrument.Metrics; import java.time.Duration; import java.util.ArrayList; import java.util.Collections; @@ -36,10 +37,12 @@ import java.util.concurrent.CompletableFuture; import java.util.function.BiFunction; import java.util.function.Supplier; +import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; public class DelayTaskRunnerTest { + private final List> runners = new ArrayList<>(); private DelayTaskRunner runner; private Supplier currentMillisSupplier; private BiFunction, CompletableFuture> taskFunction1; @@ -48,11 +51,29 @@ public class DelayTaskRunnerTest { @BeforeMethod public void setUp() { currentMillisSupplier = System::currentTimeMillis; - runner = new DelayTaskRunner<>(String::compareTo, currentMillisSupplier, 1000); + runner = newRunner(1000, "default"); taskFunction1 = mock(BiFunction.class); taskFunction2 = mock(BiFunction.class); } + /** + * Closes the runners that each test creates. + */ + @AfterMethod(alwaysRun = true) + public void tearDown() { + runners.forEach(DelayTaskRunner::shutdown); + runners.clear(); + await().atMost(Duration.ofSeconds(5)).until(() -> Metrics.globalRegistry.getMeters().stream() + .noneMatch(meter -> "DelayTaskRunnerTest".equals(meter.getId().getTag("test")))); + } + + private DelayTaskRunner newRunner(int rateLimit, String id) { + DelayTaskRunner created = new DelayTaskRunner<>(String::compareTo, currentMillisSupplier, rateLimit, + "test", "DelayTaskRunnerTest", "runner", id); + runners.add(created); + return created; + } + @Test public void testScheduleIfAbsent() { TestDelayedTask first = new TestDelayedTask(Duration.ZERO, taskFunction1); @@ -121,8 +142,7 @@ public void testHasTask() { @Test public void testRateLimiterBehavior() { int rateLimit = 2; - DelayTaskRunner rateLimitedRunner = new DelayTaskRunner<>(String::compareTo, System::currentTimeMillis, - rateLimit); + DelayTaskRunner rateLimitedRunner = newRunner(rateLimit, "rate-limited"); List executionTimes = Collections.synchronizedList(new ArrayList<>()); @@ -146,6 +166,7 @@ public void testRateLimiterBehavior() { assertTrue(diff2 >= 1800); rateLimitedRunner.shutdown(); + runners.remove(rateLimitedRunner); await().atMost(Duration.ofSeconds(5)) .until(() -> !rateLimitedRunner.hasTask("task1") && !rateLimitedRunner.hasTask("task2") && @@ -178,6 +199,7 @@ public void testShutdown() { TestDelayedTask task = new TestDelayedTask(Duration.ofMillis(50), taskFunction1); runner.schedule("shutdownKey2", task); runner.shutdown(); + runners.remove(runner); verify(taskFunction1, timeout(100).times(0)).apply("shutdownKey2", runner); }