From 683e44811fd6f7467800c0b0655dd781683d5a6e Mon Sep 17 00:00:00 2001 From: Yonny Hao Date: Mon, 26 Oct 2026 15:39:02 +0800 Subject: [PATCH 1/4] Manage executor meter ownership and shutdown cleanup --- base-util/pom.xml | 5 + .../bifromq/base/util/MonitoredExecutors.java | 431 ++++++++++++ .../util/MonitoredExecutorsExportTest.java | 128 ++++ .../base/util/MonitoredExecutorsTest.java | 652 ++++++++++++++++++ 4 files changed, 1216 insertions(+) create mode 100644 base-util/src/main/java/org/apache/bifromq/base/util/MonitoredExecutors.java create mode 100644 base-util/src/test/java/org/apache/bifromq/base/util/MonitoredExecutorsExportTest.java create mode 100644 base-util/src/test/java/org/apache/bifromq/base/util/MonitoredExecutorsTest.java 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; + } + } + +} From c01cb9d5e7a7e027b2c48dd33d0189b3699d1365 Mon Sep 17 00:00:00 2001 From: Yonny Hao Date: Mon, 26 Oct 2026 17:26:20 +0800 Subject: [PATCH 2/4] Fix executor shutdown and metrics in storage and cluster Close RocksDB and AgentHost executors and remove their meters. Remove Raft and range executor meters on shutdown. Release manager meters before range destruction completes. Keep snapshot dump metrics under a separate name. Add tests for shutdown, range split, destruction, and replacement. --- .../apache/bifromq/basecluster/AgentHost.java | 11 +- .../AgentHostExecutorLifecycleTest.java | 125 +++++++ base-kv/base-kv-local-engine-rocksdb/pom.xml | 4 + .../rocksdb/RocksDBCPableKVSpace.java | 31 +- .../localengine/rocksdb/RocksDBKVSpace.java | 13 +- .../rocksdb/RocksDBWALableKVSpace.java | 31 +- .../rocksdb/RocksDBExecutorLifecycleTest.java | 336 ++++++++++++++++++ .../apache/bifromq/basekv/raft/RaftNode.java | 4 +- .../bifromq/basekv/raft/RaftNodeTest.java | 28 +- .../store/range/KVRangeDumpSession.java | 2 +- .../basekv/store/range/KVRangeFSM.java | 10 +- .../basekv/store/KVRangeStoreTest.java | 17 + .../store/range/KVRangeDumpSessionTest.java | 80 ++++- .../KVRangeFSMExecutorLifecycleTest.java | 120 +++++++ .../basekv/store/wal/KVRangeWALTest.java | 7 +- 15 files changed, 767 insertions(+), 52 deletions(-) create mode 100644 base-cluster/src/test/java/org/apache/bifromq/basecluster/AgentHostExecutorLifecycleTest.java create mode 100644 base-kv/base-kv-local-engine-rocksdb/src/test/java/org/apache/bifromq/basekv/localengine/rocksdb/RocksDBExecutorLifecycleTest.java create mode 100644 base-kv/base-kv-store-server/src/test/java/org/apache/bifromq/basekv/store/range/KVRangeFSMExecutorLifecycleTest.java 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-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 From 343fdbe894a69cbbe44786922ff344c44a3ad4df Mon Sep 17 00:00:00 2001 From: Yonny Hao Date: Mon, 26 Oct 2026 22:30:46 +0800 Subject: [PATCH 3/4] Manage inbox and dist executor metrics per range Use range tags to give each delay and delivery executor its own meters. Remove those meters when the owner closes. Keep DelayTaskRunner.shutdown synchronous. Cancel pending triggers and wait for the executor to terminate before returning. Add tests for meter ownership and synchronous shutdown. --- bifromq-dist/bifromq-dist-worker/pom.xml | 4 + .../bifromq/dist/worker/DeliverExecutor.java | 10 +- .../dist/worker/DeliverExecutorGroup.java | 5 +- .../dist/worker/DistWorkerCoProcFactory.java | 3 +- .../worker/DeliverExecutorLifecycleTest.java | 87 ++++++++++ .../bifromq/inbox/store/InboxStoreCoProc.java | 2 +- .../inbox/store/delay/DelayTaskRunner.java | 24 ++- .../inbox/store/delay/IDelayTaskRunner.java | 2 +- .../delay/DelayTaskRunnerLifecycleTest.java | 148 ++++++++++++++++++ .../store/delay/DelayTaskRunnerTest.java | 28 +++- 10 files changed, 294 insertions(+), 19 deletions(-) create mode 100644 bifromq-dist/bifromq-dist-worker/src/test/java/org/apache/bifromq/dist/worker/DeliverExecutorLifecycleTest.java create mode 100644 bifromq-inbox/bifromq-inbox-store/src/test/java/org/apache/bifromq/inbox/store/delay/DelayTaskRunnerLifecycleTest.java 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); } From 2782d3822ccf3205f956dcb2add18a8ffbd45ea8 Mon Sep 17 00:00:00 2001 From: Yonny Hao Date: Tue, 27 Oct 2026 15:47:30 +0800 Subject: [PATCH 4/4] Close raft group after config migration test Always stop the test group so later tests can create nodes with the same metric identity. --- .../functest/RaftConfigMigrationTest.java | 48 ++++++++++--------- 1 file changed, 26 insertions(+), 22 deletions(-) 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(); + } } }