From f216c116890725b263c10a3846fd275e688342a2 Mon Sep 17 00:00:00 2001 From: Wu Sheng Date: Fri, 14 Aug 2026 22:56:58 +0800 Subject: [PATCH] Add ThreadPolicy.ioBound for blocking consumers, fix BatchQueue shutdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ioBound(N)` declares that a queue's consumers spend most of their time blocked, typically on I/O. Such a queue runs its drain loops on virtual threads where the runtime provides them (JDK 25+, via VirtualThreads.createScheduledExecutor) and on N platform threads otherwise. The count is identical on both paths, so concurrency, batching, back-pressure, drop semantics and per-partition ordering are unaffected by the fallback — only the thread substrate changes, plus shutdown latency, since the platform scheduler can drop drain tasks parked on their idle backoff and the virtual adapter cannot. The purpose is narrow: avoid paying for N parked OS threads. It does not raise concurrency, because the drain loop is the task processor and nothing is handed off — consumer concurrency is min(threadCount, partitionCount) either way, so an IO-bound queue wants threads and partitions 1:1 with a small per-partition buffer. There is deliberately no CPU-proportional form and no cpuCores(...).ioBound(): virtual threads are not preemptive, so CPU-bound work would hold its carrier and starve the pool shared with GRPCServer and HTTPServer. L1/L2/TopN stay on cpuCores/fixed. No shipped queue uses the new policy. shutdown() separately had a pre-existing race: it ran the final drain on the caller's thread while drain loops could still be running, invoking a handler from two threads and breaking the single-drain-thread invariant workers such as MetricsAggregateWorker document. It now cancels the periodic rebalance task — which does not read `running`, and whose loop on a virtual-thread scheduler exits only on interrupt — shuts the scheduler down, waits for orderly exit, and takes an exclusive dispatch lock before draining. That lock, not the wait, is the guarantee: awaitTermination is a courtesy that may time out or be interrupted. Drain loops hold the read lock for the whole cycle — the running recheck, the dequeue, notifyIdle() and the dispatch. All four are needed: an unlocked dequeue lets shutdown dispatch a newer batch ahead of one a task already holds, so the older batch lands after shutdown() returned; notifyIdle() outside the lock runs onIdle() concurrently with the final dispatch, and MetricsAggregateWorker flushes the same worker state from onIdle(); and the loop's running test can pass just as shutdown completes. Read locks are shared, so drain loops still dispatch concurrently with one another. A consumer is never interrupted mid-batch, so no truncated executeBatch or stream send. Callers that lose the shutdown CAS await the winner's completion latch, unbounded — the winner may spend all of shutdownTimeoutMs in awaitTermination and then block on the write lock for as long as a consumer runs, so any fixed bound would let a loser return while the queue is still draining. The platform scheduler also sets setExecuteExistingDelayedTasksAfterShutdownPolicy(false) so drain tasks on their idle backoff do not hold up termination; interrupts are deferred and restored only after dispatch returns. Also fixes a long-standing test bug that has kept master's CI red. Four oal-rt tests register scopes in @BeforeAll with a single try around the whole block. notify() throws when a scope id is already registered, so once any sibling has registered even one scope, the first call throws and every scope after it is silently skipped — the tests then fail at use time with "ScopeDefine name = ... not found". OALSourceAttributionTest registers Service alone, so whenever surefire's file-system order puts it first, OALClassGeneratorV2Test loses Endpoint, ServiceRelation, K8SServiceInstance and TCPService. That order occurs on the CI runners but not on macOS, which is why it reproduces only in CI. Reproduced locally with -Dsurefire.runOrder=reversealphabetical, and each registration is now guarded independently. The module passes under filesystem, alphabetical and reversealphabetical ordering. --- docs/en/changes/changes.md | 1 + .../MetricDefinitionEnricherTest.java | 28 +- .../v2/generator/OALClassGeneratorV2Test.java | 28 +- .../library-batch-queue/CLAUDE.md | 124 +++++- .../library-batch-queue/pom.xml | 5 + .../server/library/batchqueue/BatchQueue.java | 253 ++++++++++-- .../library/batchqueue/BatchQueueConfig.java | 15 + .../library/batchqueue/ThreadPolicy.java | 59 ++- .../batchqueue/BatchQueueShutdownTest.java | 361 ++++++++++++++++++ .../library/batchqueue/ThreadPolicyTest.java | 29 ++ 10 files changed, 832 insertions(+), 71 deletions(-) create mode 100644 oap-server/server-library/library-batch-queue/src/test/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueueShutdownTest.java diff --git a/docs/en/changes/changes.md b/docs/en/changes/changes.md index 455b12a197d4..3bfcb36184fc 100644 --- a/docs/en/changes/changes.md +++ b/docs/en/changes/changes.md @@ -256,6 +256,7 @@ admin-host only" entry above for the public REST retirement. * Drop six unused test-scoped dependencies from `runtime-rule` (`library-integration-test`, `library-banyandb-client`, `storage-banyandb-plugin`, `testcontainers`, `testcontainers:junit-jupiter`, `grpc-testing`). They staged the plugin-side ITs that were retired in favour of e2e; that coverage now lives in `test/e2e-v2/cases/runtime-rule/` (MAL over BanyanDB / PostgreSQL / Elasticsearch, LAL, meter, and the two-node cluster case). The module has no ITs today, and JUnit and Mockito are inherited from the root POM. * Declare `server-testing` at `test` scope everywhere. It ships only test scaffolding (`ModuleManagerTesting`, `MockModuleManager`, the MAL/LAL/Hierarchy rule loaders) plus two empty `org.junit` stubs that let Testcontainers' `GenericContainer` hierarchy resolve without JUnit 4, but four modules declared it at compile scope — including the `server-configuration` parent, so all eight `configuration-*` children inherited it — which put those `org.junit` stubs on the runtime classpath that `server-starter` copies into `oap-libs`. Modules whose tests need the stubs now declare the dependency themselves rather than inheriting it transitively, and `library-banyandb-client` gains the direct `library-util` dependency its `BanyanDBClient` always needed (it was resolving `StringUtil` through `server-testing`, a test-support module). +* Add `ThreadPolicy.ioBound(N)` to `library-batch-queue`, for queues whose consumers spend most of their time blocked. Such a queue runs its drain loops on virtual threads where the runtime provides them (JDK 25+) and falls back to N platform threads otherwise; the count, and therefore concurrency, batching, back-pressure, drop semantics and per-partition ordering, are identical on both paths. Shutdown latency is the one exception: the platform scheduler drops drain tasks parked on their idle backoff, while the virtual-thread adapter sleeps inside the submitted task and cannot, so an `ioBound` queue should keep `maxIdleMs` within `shutdownTimeoutMs`. There is deliberately no CPU-proportional form: virtual threads are not preemptive, so CPU-bound work would hold its carrier and starve the shared carrier pool, and L1/L2/TopN stay on `cpuCores`/`fixed`. Also fixes `BatchQueue.shutdown()`, which ran its final drain on the caller's thread while drain loops could still be inside `consume()`, invoking a handler concurrently and breaking the single-drain-thread invariant workers such as `MetricsAggregateWorker` rely on: it now cancels the periodic rebalance task, waits for in-flight consumers (`shutdownTimeoutMs`, default 500ms per queue), and serialises its final dispatch behind a read/write dispatch lock so the guarantee holds even when that wait times out or is interrupted. Drain loops hold the read lock for the whole cycle — the running recheck, the partition dequeue, the idle notification and the dispatch — because `onIdle()` touches the same worker state as `consume()` and an unlocked dequeue would let shutdown dispatch a newer batch ahead of one a task already holds. Concurrent shutdown callers await the winner's completion rather than returning early. A consumer is never interrupted mid-batch. #### OAP Server * Fix LAL's `segmentId` and `spanId` extractor statements, which the grammar accepted and the parser never implemented. `LALParser.g4` declares `traceIdStatement`, `segmentIdStatement` and `spanIdStatement`, and the codegen already carried `setSegmentId`/`setSpanId` in its setter table, but `LALScriptParser.visitExtractorStatement` had a branch for only the first of the three. The remaining alternatives fell through to a line that assumed whatever was left had to be an `ifStatement`, so a rule writing `segmentId ...` failed at boot with a `NullPointerException` naming `IfStatementContext` — for a rule line containing no `if`. Both statements now work, and an unhandled extractor statement reports its own rule line instead of throwing. Existing log records are unaffected: `LogBuilder` copies trace id, segment id and span id straight from the log's metadata, and only skips that copy when a rule has set them — which no shipped rule did, which is why the gap went unnoticed. Dedicated execution tests now cover reading all three fields from `log.traceContext.*` and writing all three from an extractor. diff --git a/oap-server/oal-rt/src/test/java/org/apache/skywalking/oal/v2/generator/MetricDefinitionEnricherTest.java b/oap-server/oal-rt/src/test/java/org/apache/skywalking/oal/v2/generator/MetricDefinitionEnricherTest.java index 9b59a8cf6291..b79a6ecc00fe 100644 --- a/oap-server/oal-rt/src/test/java/org/apache/skywalking/oal/v2/generator/MetricDefinitionEnricherTest.java +++ b/oap-server/oal-rt/src/test/java/org/apache/skywalking/oal/v2/generator/MetricDefinitionEnricherTest.java @@ -53,16 +53,24 @@ public class MetricDefinitionEnricherTest { @BeforeAll public static void initializeScopes() { - try { - DefaultScopeDefine.Listener listener = new DefaultScopeDefine.Listener(); - listener.notify(Service.class); - listener.notify(Endpoint.class); - listener.notify(ServiceRelation.class); - listener.notify(K8SService.class); - listener.notify(K8SServiceInstance.class); - listener.notify(TCPService.class); - } catch (RuntimeException e) { - // Scopes may already be registered by other tests + // Each registration is guarded separately. notify() throws when a scope id is already + // registered by another test in this JVM, and a single try around the whole block would + // swallow that and silently skip every scope after it — leaving the rest unregistered and + // failing later with "ScopeDefine name = ... not found". Which scope trips first depends + // on surefire's run order, so the truncation only shows up on some platforms. + final DefaultScopeDefine.Listener listener = new DefaultScopeDefine.Listener(); + for (final Class scope : List.of( + Service.class, + Endpoint.class, + ServiceRelation.class, + K8SService.class, + K8SServiceInstance.class, + TCPService.class)) { + try { + listener.notify(scope); + } catch (RuntimeException e) { + // Already registered by a sibling test; the remaining scopes still need registering. + } } } diff --git a/oap-server/oal-rt/src/test/java/org/apache/skywalking/oal/v2/generator/OALClassGeneratorV2Test.java b/oap-server/oal-rt/src/test/java/org/apache/skywalking/oal/v2/generator/OALClassGeneratorV2Test.java index d3adb48863e3..5814039d987a 100644 --- a/oap-server/oal-rt/src/test/java/org/apache/skywalking/oal/v2/generator/OALClassGeneratorV2Test.java +++ b/oap-server/oal-rt/src/test/java/org/apache/skywalking/oal/v2/generator/OALClassGeneratorV2Test.java @@ -72,16 +72,24 @@ public class OALClassGeneratorV2Test { @BeforeAll public static void initializeScopes() { - try { - DefaultScopeDefine.Listener listener = new DefaultScopeDefine.Listener(); - listener.notify(Service.class); - listener.notify(Endpoint.class); - listener.notify(ServiceRelation.class); - listener.notify(K8SService.class); - listener.notify(K8SServiceInstance.class); - listener.notify(TCPService.class); - } catch (RuntimeException e) { - // Scopes may already be registered by other tests + // Each registration is guarded separately. notify() throws when a scope id is already + // registered by another test in this JVM, and a single try around the whole block would + // swallow that and silently skip every scope after it — leaving the rest unregistered and + // failing later with "ScopeDefine name = ... not found". Which scope trips first depends + // on surefire's run order, so the truncation only shows up on some platforms. + final DefaultScopeDefine.Listener listener = new DefaultScopeDefine.Listener(); + for (final Class scope : List.of( + Service.class, + Endpoint.class, + ServiceRelation.class, + K8SService.class, + K8SServiceInstance.class, + TCPService.class)) { + try { + listener.notify(scope); + } catch (RuntimeException e) { + // Already registered by a sibling test; the remaining scopes still need registering. + } } } diff --git a/oap-server/server-library/library-batch-queue/CLAUDE.md b/oap-server/server-library/library-batch-queue/CLAUDE.md index ad88ca6ba350..f27a0dd49c84 100644 --- a/oap-server/server-library/library-batch-queue/CLAUDE.md +++ b/oap-server/server-library/library-batch-queue/CLAUDE.md @@ -48,7 +48,7 @@ use `BatchQueueManager.create(name, config)` which throws on duplicate names. | `BatchQueue` | The queue itself. Holds partitions, runs drain loops, dispatches to consumers/handlers. | | `BatchQueueManager` | Global registry. Creates/retrieves queues by name. `create()` for unique, `getOrCreate()` for shared. | | `BatchQueueConfig` | Builder for queue configuration (threads, partitions, buffer, strategy, consumer, balancer). | -| `ThreadPolicy` | Resolves thread count: `fixed(N)`, `cpuCores(mult)`, `cpuCoresWithBase(base, mult)`. | +| `ThreadPolicy` | Resolves thread count: `fixed(N)`, `cpuCores(mult)`, `cpuCoresWithBase(base, mult)`, `ioBound(N)`. | | `PartitionPolicy` | Resolves partition count: `fixed(N)`, `threadMultiply(N)`, `adaptive()`. | | `PartitionSelector` | Routes items to partitions. Default `typeHash()` groups by class. | | `HandlerConsumer` | Callback for processing a batch. Has optional `onIdle()` for flush-on-idle. | @@ -60,13 +60,59 @@ use `BatchQueueManager.create(name, config)` which throws on duplicate names. ## ThreadPolicy ```java -ThreadPolicy.fixed(4) // exactly 4 threads -ThreadPolicy.cpuCores(1.0) // 1 thread per CPU core +ThreadPolicy.fixed(4) // exactly 4 threads +ThreadPolicy.cpuCores(1.0) // 1 thread per CPU core ThreadPolicy.cpuCoresWithBase(1, 0.25) // 1 + 0.25 * cores (e.g., 3 on 8-core) +ThreadPolicy.ioBound(50) // 50 threads whose consumers block; virtual when available ``` Always resolves to >= 1. +### ioBound — blocking consumers + +`ioBound(n)` declares that this queue's consumers spend most of their time blocked, typically on +I/O. The scheduler then uses virtual threads where the runtime provides them (JDK 25+, via +`VirtualThreads.createScheduledExecutor`) and falls back to `n` platform threads otherwise. + +**The count is identical on both paths.** Concurrency, batching, back-pressure, drop semantics and +per-partition ordering are unaffected by the fallback. The queue logs once at creation when it +asked for virtual threads and got platform ones. + +One behaviour does differ: **shutdown latency**. On the platform path the scheduler drops drain +tasks still waiting out their idle backoff, so termination is immediate. The virtual-thread adapter +implements a delay by sleeping *inside* the submitted task, which `shutdown()` cannot cancel, so +termination waits out the longest outstanding backoff — bounded by `maxIdleMs`. Keep +`maxIdleMs <= shutdownTimeoutMs` on an `ioBound` queue, or `shutdown()` will log a spurious +"drain loops did not finish" warning on every teardown. Nothing is lost either way: a late-waking +drain task sees `running == false` and exits without draining, and the final drain flushes the data. + +Sizing: the drain loop **is** the task processor — a consumer blocking for seconds occupies its +loop for that whole time, and nothing is handed off. So consumer concurrency is +`min(threadCount, partitionCount)`, and an IO-bound queue wants **threads and partitions 1:1** +with a small per-partition buffer: + +```java +.threads(ThreadPolicy.ioBound(concurrency)) +.partitions(PartitionPolicy.fixed(concurrency)) // 1:1, or the count is clamped away +.bufferSize(4) // per partition +.strategy(BufferStrategy.IF_POSSIBLE) +.minIdleMs(50).maxIdleMs(500) // seconds-long work; 1ms polling is waste +``` + +The count expresses **how many concurrent blocking calls the downstream service tolerates**, which +does not follow core count — hence no CPU-proportional variant, and hence the count is mandatory. + +**Never use `ioBound` for CPU-bound work.** Virtual threads are not preemptive: a CPU-only task +holds its carrier to completion, and the carrier pool is shared process-wide with the executors +`GRPCServer` and `HTTPServer` already create. L1 (`MetricsAggregateWorker`), L2 +(`MetricsPersistentMinWorker`) and TopN do in-memory merges across hundreds of partitions and must +stay on `cpuCores`/`fixed`. The API has no `cpuCores(...).ioBound()` form so this combination +cannot be expressed. + +`ioBound` with a `DrainBalancer` logs a warning: rebalancing exists to redistribute skewed +CPU-bound partitions, which is not this shape. It is not rejected — it remains safe — but it is +usually a configuration mistake. + ## PartitionPolicy ```java @@ -221,5 +267,73 @@ mode: single consumer (JDBC batch flush) 3. `queue.addHandler(type, handler)` -- registers type handler (adaptive: may grow partitions) 4. `queue.produce(data)` -- routes to partition, blocks or drops per strategy 5. Drain loops run continuously, dispatching batches to consumers/handlers -6. `BatchQueueManager.shutdown(name)` -- stops drain, final flush -7. `BatchQueueManager.shutdownAll()` -- called during OAP server shutdown +6. `BatchQueueManager.shutdown(name)` -- stops drain, waits for in-flight consumers, final flush +7. `BatchQueueManager.shutdownAll()` -- available, but currently has no production callers + +### Shutdown order + +``` +CAS shutdownStarted winner runs the sequence; losers await its completion latch +running = false reject produce(); drain chains stop re-queueing +rebalanceFuture.cancel(true) the periodic task does not read `running` +scheduler.shutdown() +awaitTermination(config) COURTESY wait for orderly exit; may time out or be interrupted +dispatchLock.writeLock() THIS is the guarantee — waits out any in-flight consume() +final drain + dispatch +restore interrupt only if interrupted, and only after dispatch +``` + +The order is load-bearing: + +- **cancel before awaitTermination** -- on a virtual-thread scheduler the periodic task is one + submission whose loop exits only on interrupt, and `shutdown()` does not interrupt. Without the + cancel, termination can never be observed. (A `ScheduledThreadPoolExecutor` cancels periodic + tasks itself, so the platform path worked by relying on that default.) +- **`running = false` before cancel** -- the rebalance fence spins on + `cycleCount.get(t) <= snap && running`, and `LockSupport.parkNanos` *returns* on interrupt + without throwing, leaving the flag set. The fence exits via `running`, not the interrupt. +- **wait before the final drain** -- draining on the caller's thread while a drain loop is still + inside `consume()` invokes the same handler concurrently, breaking the single-drain-thread + invariant workers such as `MetricsAggregateWorker` depend on. + +The final drain deliberately ignores partition ownership, so partitions left `UNOWNED` by an +interrupted rebalance are still flushed. + +**`awaitTermination` is not the safety mechanism -- the dispatch lock is.** The wait is a courtesy +for orderly exit and can time out or be interrupted, and on either path a drain loop may still be +running. `shutdown()` takes the WRITE lock, so its final dispatch waits the in-flight consumer out +instead of racing it. + +Drain loops hold the READ lock for the **whole cycle** -- the `running` recheck, the partition +dequeue, `notifyIdle()` and the dispatch -- not merely around `dispatch()`. All four are required: +a dequeue outside the lock lets shutdown drain and dispatch a *newer* batch while a task holds an +older one, which is then dispatched out of order after `shutdown()` has returned; `notifyIdle()` +outside it runs `onIdle()` concurrently with the final dispatch, and implementations such as +`MetricsAggregateWorker` flush the same worker state from `onIdle()`; and the loop's `running` test +can pass just as shutdown completes. Read locks are shared, so drain loops still run concurrently +with one another -- correct, since the partition selector routes a type to one partition and +therefore one task. Nothing sleeps inside the body (backoff is applied by `scheduleDrain` after the +loop exits), so the lock is held only for as long as the consumer runs. Blocking there is bounded by that consumer's +own work and is preferable to interrupting it: no consumer is written to tolerate interruption +mid-batch, and a truncated `executeBatch` or stream send is worse than the race being fixed, which +is why `shutdownNow()` is never called. + +`shutdownTimeoutMs` is per-queue (default 500ms) and bounds only the courtesy wait; `0` skips it +entirely, which is safe now that the lock -- not the wait -- provides the guarantee. + +Two further properties, both for the same reason -- a consumer must never be entered twice at once, +nor with the interrupt flag set: + +- **Once-only, and shared completion.** Only the caller that wins a CAS runs the sequence; the + losers block on a completion latch rather than returning early, so `shutdownAll()` racing a + `shutdown(name)` cannot report completion while the winner is still draining. That wait is + deliberately **unbounded**: the winner may spend all of `shutdownTimeoutMs` in `awaitTermination` + and then block on the write lock for as long as a consumer runs, so any fixed bound would let a + loser return early. Interruption is deferred and restored on the way out. +- **Interrupt is deferred.** If `awaitTermination` is interrupted, the queue logs, still runs the + final drain (skipping it would lose data), and re-asserts the thread's interrupt flag only after + `dispatch()` returns. + +The platform scheduler additionally sets `setExecuteExistingDelayedTasksAfterShutdownPolicy(false)` +so drain tasks parked on their idle backoff do not hold up termination. Dropping them is safe: such +a task, had it run, would find `running == false` and exit at the top of its loop without draining. diff --git a/oap-server/server-library/library-batch-queue/pom.xml b/oap-server/server-library/library-batch-queue/pom.xml index 455b77f4478a..eb300f08cf12 100644 --- a/oap-server/server-library/library-batch-queue/pom.xml +++ b/oap-server/server-library/library-batch-queue/pom.xml @@ -28,6 +28,11 @@ library-batch-queue + + org.apache.skywalking + library-util + ${project.version} + org.awaitility awaitility diff --git a/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueue.java b/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueue.java index 1890bd0d2bfe..8fc1e9b92c46 100644 --- a/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueue.java +++ b/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueue.java @@ -26,15 +26,21 @@ import java.util.Set; import java.util.concurrent.ArrayBlockingQueue; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicIntegerArray; import java.util.concurrent.atomic.AtomicLongArray; import java.util.concurrent.locks.LockSupport; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.function.Supplier; import lombok.AccessLevel; import lombok.Getter; import lombok.extern.slf4j.Slf4j; +import org.apache.skywalking.oap.server.library.util.VirtualThreads; /** * A partitioned, self-draining queue with type-based dispatch. @@ -114,6 +120,42 @@ public class BatchQueue { /** The thread pool that executes drain tasks. */ private final ScheduledExecutorService scheduler; + /** + * Handle on the periodic rebalance task, so {@link #shutdown()} can stop it explicitly. + * A {@code ScheduledThreadPoolExecutor} cancels periodic tasks on {@code shutdown()} by + * default, but a virtual-thread-backed scheduler runs the task as one submission whose loop + * exits only on interrupt — and {@code shutdown()} does not interrupt. Without this handle + * {@code awaitTermination} could never succeed on that path. + */ + private volatile ScheduledFuture rebalanceFuture; + + /** + * Guards {@link #shutdown()} so at most one caller ever runs the final drain. Without it two + * concurrent shutdowns — {@code shutdownAll()} racing {@code shutdown(name)}, since the former + * snapshots the registry before clearing it — could interleave inside the final-drain loop, + * each taking some partitions and dispatching, invoking the same consumer from two threads. + * + *

Defensive: the interleaving window is inside the per-partition loop and is too narrow to + * pin down in a deterministic test, so this guard is not covered by one. + */ + private final AtomicBoolean shutdownStarted = new AtomicBoolean(); + + /** Released by the shutdown winner so losing callers do not return before it has finished. */ + private final CountDownLatch shutdownComplete = new CountDownLatch(1); + + /** + * Serialises the shutdown-time dispatch against in-flight drain dispatches. + * + *

Drain loops take the READ lock, so they still dispatch concurrently with one another — + * correct, because the partition selector routes a given type to one partition and therefore + * one task. {@link #shutdown()} takes the WRITE lock, so its final dispatch waits for every + * in-flight {@code consume()} to return instead of racing it. This is what makes the guarantee + * hold even when {@code awaitTermination} times out or is interrupted: waiting politely for + * orderly termination is bounded, but never invoking a consumer from two threads is not + * negotiable. + */ + private final ReentrantReadWriteLock dispatchLock = new ReentrantReadWriteLock(); + /** * Cached partition selector from config. Only used when {@code partitions.length > 1}; * single-partition queues bypass the selector entirely. @@ -329,12 +371,39 @@ public class BatchQueue { partitions[i] = new ArrayBlockingQueue<>(config.getBufferSize()); } - this.scheduler = Executors.newScheduledThreadPool(threadCount, r -> { - final Thread t = new Thread(r); - t.setName("BatchQueue-" + name + "-" + t.getId()); - t.setDaemon(true); - return t; - }); + // threadCount is reassigned by the partition clamp above; copy for the lambda. + final int poolSize = threadCount; + final Supplier platformScheduler = () -> { + final ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(poolSize, r -> { + final Thread t = new Thread(r); + t.setName("BatchQueue-" + name + "-" + t.getId()); + t.setDaemon(true); + return t; + }); + // Drop drain tasks still waiting out their idle backoff when shutdown() runs. The + // default keeps them queued, so a queue with a long maxIdleMs could not terminate + // until the delay elapsed. Dropping them is safe: a drain task that did run after + // `running` went false exits at the top of its loop without draining, and the final + // drain flushes the data regardless. + executor.setExecuteExistingDelayedTasksAfterShutdownPolicy(false); + return executor; + }; + if (config.getThreads().isIoBound()) { + // Virtual threads when the runtime has them, otherwise poolSize platform threads. + // The count is the same either way, so only the substrate differs. + this.scheduler = VirtualThreads.createScheduledExecutor("BatchQueue-" + name, platformScheduler); + if (!VirtualThreads.isSupported()) { + log.info("BatchQueue[{}]: IO-bound policy requested but virtual threads are " + + "unavailable, running {} platform threads instead", name, poolSize); + } + if (config.getBalancer() != null) { + log.warn("BatchQueue[{}]: IO-bound policy combined with drain rebalancing — " + + "rebalancing targets skewed CPU-bound partitions, which is not the " + + "blocking workload this policy is for", name); + } + } else { + this.scheduler = platformScheduler.get(); + } this.taskCount = threadCount; this.assignedPartitions = buildAssignments(threadCount, partitionCount); @@ -539,8 +608,9 @@ private void enableRebalancing(final long intervalMs) { // Enable the flag — gates hot-path additions in produce() and drainLoop() this.rebalancingEnabled = true; - // Schedule periodic rebalancing on the queue's scheduler - scheduler.scheduleAtFixedRate( + // Schedule periodic rebalancing on the queue's scheduler. The handle is kept so + // shutdown() can cancel it — see the rebalanceFuture field. + this.rebalanceFuture = scheduler.scheduleAtFixedRate( this::rebalance, intervalMs, intervalMs, TimeUnit.MILLISECONDS ); @@ -608,28 +678,48 @@ void drainLoop(final int taskIndex) { final boolean checkOwnership = rebalancingEnabled; try { while (running) { - // Drain all assigned partitions into one batch - final List combined = new ArrayList<>(); - for (final int partitionIndex : myPartitions) { - if (partitionIndex < currentPartitions.length) { - // Skip partitions revoked by the rebalancer - if (checkOwnership && partitionOwner.get(partitionIndex) != taskIndex) { - continue; + // The read lock spans the WHOLE cycle, not just dispatch. Holding it only around + // dispatch would leave three gaps against shutdown's exclusive final drain: + // - dequeue outside it lets shutdown drain and dispatch a NEWER batch while this + // task holds an older one, which would then be dispatched out of order, after + // shutdown() has already returned; + // - notifyIdle() outside it runs onIdle() concurrently with the final dispatch, + // and implementations such as MetricsAggregateWorker flush the same worker + // state from onIdle(); + // - the `running` test above can pass just as shutdown completes. + // There is no sleep in this body — backoff is applied by scheduleDrain after the + // loop exits — so the lock is held only for as long as the consumer runs. + dispatchLock.readLock().lock(); + try { + if (!running) { + break; + } + + // Drain all assigned partitions into one batch + final List combined = new ArrayList<>(); + for (final int partitionIndex : myPartitions) { + if (partitionIndex < currentPartitions.length) { + // Skip partitions revoked by the rebalancer + if (checkOwnership && partitionOwner.get(partitionIndex) != taskIndex) { + continue; + } + currentPartitions[partitionIndex].drainTo(combined); } - currentPartitions[partitionIndex].drainTo(combined); } - } - if (combined.isEmpty()) { - // Nothing to drain — increase backoff and notify idle - consecutiveIdleCycles[taskIndex]++; - notifyIdle(taskIndex, myPartitions); - break; - } + if (combined.isEmpty()) { + // Nothing to drain — increase backoff and notify idle + consecutiveIdleCycles[taskIndex]++; + notifyIdle(taskIndex, myPartitions); + break; + } - // Data found — reset backoff and dispatch - consecutiveIdleCycles[taskIndex] = 0; - dispatch(combined, taskIndex, myPartitions); + // Data found — reset backoff and dispatch + consecutiveIdleCycles[taskIndex] = 0; + dispatch(combined, taskIndex, myPartitions); + } finally { + dispatchLock.readLock().unlock(); + } } } catch (final Throwable t) { log.error("BatchQueue[{}]: drain loop error", name, t); @@ -913,22 +1003,109 @@ private static int[][] buildAssignmentsFromOwner( } /** - * Stop the queue: reject new produces, perform a final drain of all partitions, - * and shut down the scheduler. + * Stop the queue: reject new produces, stop the drain loops, wait for in-flight consumer + * invocations to finish, then perform a final drain of all partitions. + * + *

Step order is load-bearing. {@code running = false} stops the drain chains but is not + * read by the periodic rebalance task, so that is cancelled explicitly — and it must be + * cancelled before {@code awaitTermination}, which otherwise could never observe termination. + * The rebalance fence itself exits via {@code running}, not the interrupt, because + * {@code LockSupport.parkNanos} returns on interrupt without throwing. + * + *

Waiting before the final drain is what keeps consumers single-threaded: draining on the + * caller's thread while a drain loop is still inside {@code consume()} would invoke the same + * handler concurrently. */ void shutdown() { - running = false; - // Final drain — flush any remaining data to consumers - final ArrayBlockingQueue[] currentPartitions = this.partitions; - final List combined = new ArrayList<>(); - for (final ArrayBlockingQueue partition : currentPartitions) { - partition.drainTo(combined); + if (!shutdownStarted.compareAndSet(false, true)) { + // A losing caller must not report completion before the winner has finished draining, + // or shutdownAll() can return while a consumer is still running. + awaitShutdownComplete(); + return; + } + try { + doShutdown(); + } finally { + shutdownComplete.countDown(); } - if (!combined.isEmpty()) { - // Shutdown dispatch — no idle notification needed - dispatch(combined, -1, null); + } + + /** Test hook: true once a caller has won the shutdown CAS. */ + boolean isShutdownStarted() { + return shutdownStarted.get(); + } + + /** + * Await the winner's completion without a deadline. A bounded wait cannot be correct here: the + * winner may spend the whole of {@code shutdownTimeoutMs} in {@code awaitTermination} and then + * block on the write lock for as long as an in-flight consumer runs, so any fixed bound would + * let this caller return while the queue is still draining — the very thing the shared + * completion exists to prevent. Interruption is deferred and restored on the way out. + */ + private void awaitShutdownComplete() { + boolean interrupted = false; + while (true) { + try { + shutdownComplete.await(); + break; + } catch (final InterruptedException e) { + interrupted = true; + } + } + if (interrupted) { + Thread.currentThread().interrupt(); } + } + + private void doShutdown() { + running = false; + + final ScheduledFuture currentRebalance = this.rebalanceFuture; + if (currentRebalance != null) { + currentRebalance.cancel(true); + } + scheduler.shutdown(); + final long timeoutMs = config.getShutdownTimeoutMs(); + boolean interrupted = false; + try { + if (timeoutMs > 0 && !scheduler.awaitTermination(timeoutMs, TimeUnit.MILLISECONDS)) { + log.warn("BatchQueue[{}]: drain loops did not finish within {}ms; the final drain " + + "waits for the exclusive dispatch lock rather than racing them", + name, timeoutMs); + } + } catch (final InterruptedException e) { + interrupted = true; + log.warn("BatchQueue[{}]: interrupted while waiting for drain loops; the final drain " + + "waits for the exclusive dispatch lock rather than racing them", name); + } + + // The write lock is what actually guarantees the consumer is never entered twice at once. + // awaitTermination above is only a courtesy wait for orderly exit — it may time out or be + // interrupted, and on either path a drain loop can still be inside consume(). Blocking here + // is bounded by that consumer's own work, and is preferable to interrupting it mid-batch. + dispatchLock.writeLock().lock(); + try { + // Final drain — flush any remaining data to consumers. Deliberately ignores partition + // ownership so that partitions left UNOWNED by an interrupted rebalance are drained too. + final ArrayBlockingQueue[] currentPartitions = this.partitions; + final List combined = new ArrayList<>(); + for (final ArrayBlockingQueue partition : currentPartitions) { + partition.drainTo(combined); + } + if (!combined.isEmpty()) { + // Shutdown dispatch — no idle notification needed + dispatch(combined, -1, null); + } + } finally { + dispatchLock.writeLock().unlock(); + } + + // Restored only now: consumers must not be entered with the interrupt flag set, since + // none of them is written to tolerate interruption mid-batch. + if (interrupted) { + Thread.currentThread().interrupt(); + } } int getPartitionCount() { diff --git a/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueueConfig.java b/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueueConfig.java index b47e154acb80..9cbdea97b8a9 100644 --- a/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueueConfig.java +++ b/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueueConfig.java @@ -60,6 +60,18 @@ public class BatchQueueConfig { @Builder.Default private long maxIdleMs = 50; + /** + * How long {@code shutdown()} waits for in-flight consumer invocations to finish before + * performing the final drain. Per-queue because the bound is set by the consumer's work: + * a gRPC peer teardown runs on the cluster topology thread and cannot wait long, while a + * queue whose consumer makes a multi-second remote call needs far more. + * + * On expiry the queue logs and proceeds with the final drain rather than interrupting the + * consumer — no consumer is written to tolerate interruption mid-batch. + */ + @Builder.Default + private long shutdownTimeoutMs = 500; + /** * Drain balancer for periodic rebalancing of partition-to-thread assignments. * Set via {@code .balancer(DrainBalancer, intervalMs)} on the builder. @@ -89,6 +101,9 @@ void validate() { throw new IllegalArgumentException( "maxIdleMs must be >= minIdleMs, got maxIdleMs=" + maxIdleMs + " minIdleMs=" + minIdleMs); } + if (shutdownTimeoutMs < 0) { + throw new IllegalArgumentException("shutdownTimeoutMs must be >= 0, got: " + shutdownTimeoutMs); + } } /** diff --git a/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/ThreadPolicy.java b/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/ThreadPolicy.java index b1d50f4dda9f..c71f20606046 100644 --- a/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/ThreadPolicy.java +++ b/oap-server/server-library/library-batch-queue/src/main/java/org/apache/skywalking/oap/server/library/batchqueue/ThreadPolicy.java @@ -19,25 +19,35 @@ package org.apache.skywalking.oap.server.library.batchqueue; /** - * Determines the number of threads for a BatchQueue's scheduler. + * Determines the number of threads for a BatchQueue's scheduler, and whether that queue's + * work is blocking-dominated. * - * Three modes: + * Four modes: * - fixed(N): exactly N threads, regardless of hardware. * - cpuCores(multiplier): multiplier * Runtime.availableProcessors(), rounded. * - cpuCoresWithBase(base, multiplier): base + multiplier * Runtime.availableProcessors(), rounded. + * - ioBound(N): exactly N threads whose consumers spend most of their time blocked. * * Resolved value is always >= 1 — every pool must have at least one thread. * fixed() requires count >= 1 at construction. cpuCores() applies max(1, ...) at resolution. + * + * There is deliberately no CPU-proportional {@code ioBound} variant: sizing by core count is + * meaningless for work that blocks, and virtual threads must never carry CPU-bound work — they + * are not preemptive, so a CPU-only task holds its carrier to completion and starves every other + * virtual thread in the process. */ public class ThreadPolicy { private final int fixedCount; private final int base; private final double cpuMultiplier; + private final boolean ioBound; - private ThreadPolicy(final int fixedCount, final int base, final double cpuMultiplier) { + private ThreadPolicy(final int fixedCount, final int base, final double cpuMultiplier, + final boolean ioBound) { this.fixedCount = fixedCount; this.base = base; this.cpuMultiplier = cpuMultiplier; + this.ioBound = ioBound; } /** @@ -51,7 +61,7 @@ public static ThreadPolicy fixed(final int count) { if (count < 1) { throw new IllegalArgumentException("Thread count must be >= 1, got: " + count); } - return new ThreadPolicy(count, 0, 0); + return new ThreadPolicy(count, 0, 0, false); } /** @@ -66,7 +76,7 @@ public static ThreadPolicy cpuCores(final double multiplier) { if (multiplier <= 0) { throw new IllegalArgumentException("CPU multiplier must be > 0, got: " + multiplier); } - return new ThreadPolicy(0, 0, multiplier); + return new ThreadPolicy(0, 0, multiplier, false); } /** @@ -87,7 +97,30 @@ public static ThreadPolicy cpuCoresWithBase(final int base, final double multipl if (multiplier <= 0) { throw new IllegalArgumentException("CPU multiplier must be > 0, got: " + multiplier); } - return new ThreadPolicy(0, base, multiplier); + return new ThreadPolicy(0, base, multiplier, false); + } + + /** + * Exactly {@code count} threads for a queue whose consumers spend most of their time blocked — + * typically on I/O. + * + * Virtual threads are used when the runtime supports them (JDK 25+, see + * {@code VirtualThreads}); otherwise the queue falls back to {@code count} platform threads. + * The count is identical either way — only the thread substrate changes, so concurrency, + * batching, back-pressure and per-partition ordering are unaffected by the fallback. + * + * The count is mandatory and has no CPU-proportional form: it expresses how many concurrent + * blocking calls the downstream service tolerates, which does not follow core count. + * + * @param count the exact number of threads + * @return a ThreadPolicy with a fixed thread count, marked as blocking-dominated + * @throws IllegalArgumentException if count < 1 + */ + public static ThreadPolicy ioBound(final int count) { + if (count < 1) { + throw new IllegalArgumentException("Thread count must be >= 1, got: " + count); + } + return new ThreadPolicy(count, 0, 0, true); } /** @@ -102,6 +135,14 @@ public int resolve() { return Math.max(1, base + (int) Math.round(cpuMultiplier * Runtime.getRuntime().availableProcessors())); } + /** + * @return true when this queue's consumers are expected to block, making virtual threads + * appropriate where the runtime provides them + */ + boolean isIoBound() { + return ioBound; + } + @Override public boolean equals(final Object o) { if (this == o) { @@ -113,7 +154,8 @@ public boolean equals(final Object o) { final ThreadPolicy that = (ThreadPolicy) o; return fixedCount == that.fixedCount && base == that.base - && Double.compare(that.cpuMultiplier, cpuMultiplier) == 0; + && Double.compare(that.cpuMultiplier, cpuMultiplier) == 0 + && ioBound == that.ioBound; } @Override @@ -122,13 +164,14 @@ public int hashCode() { result = 31 * result + base; final long temp = Double.doubleToLongBits(cpuMultiplier); result = 31 * result + (int) (temp ^ (temp >>> 32)); + result = 31 * result + (ioBound ? 1 : 0); return result; } @Override public String toString() { if (fixedCount > 0) { - return "fixed(" + fixedCount + ")"; + return ioBound ? "ioBound(" + fixedCount + ")" : "fixed(" + fixedCount + ")"; } if (base > 0) { return "cpuCoresWithBase(" + base + ", " + cpuMultiplier + ")"; diff --git a/oap-server/server-library/library-batch-queue/src/test/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueueShutdownTest.java b/oap-server/server-library/library-batch-queue/src/test/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueueShutdownTest.java new file mode 100644 index 000000000000..1711e71d515e --- /dev/null +++ b/oap-server/server-library/library-batch-queue/src/test/java/org/apache/skywalking/oap/server/library/batchqueue/BatchQueueShutdownTest.java @@ -0,0 +1,361 @@ +/* + * 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.skywalking.oap.server.library.batchqueue; + +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.awaitility.Awaitility; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class BatchQueueShutdownTest { + + @AfterEach + public void cleanup() { + BatchQueueManager.reset(); + } + + /** + * shutdown() must not run the final drain while a drain loop is still inside consume(). + * Before the fix the caller dispatched immediately, so the handler saw two concurrent + * invocations — the invariant workers such as MetricsAggregateWorker rely on. + */ + @Test + public void testShutdownWaitsForInFlightConsumer() throws Exception { + final CountDownLatch consumerEntered = new CountDownLatch(1); + final AtomicInteger concurrentConsumers = new AtomicInteger(); + final AtomicInteger maxConcurrentConsumers = new AtomicInteger(); + + final BatchQueue queue = BatchQueueManager.create("shutdown-wait-test", + BatchQueueConfig.builder() + .threads(ThreadPolicy.fixed(1)) + .partitions(PartitionPolicy.fixed(1)) + .bufferSize(100) + .shutdownTimeoutMs(5_000) + .consumer(batch -> { + maxConcurrentConsumers.accumulateAndGet( + concurrentConsumers.incrementAndGet(), Math::max); + try { + consumerEntered.countDown(); + Thread.sleep(300); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + concurrentConsumers.decrementAndGet(); + } + }) + .build()); + + queue.produce("first"); + assertTrue(consumerEntered.await(5, TimeUnit.SECONDS), "consumer never started"); + // Queued behind the consumer that is currently sleeping, so it is still buffered + // when shutdown() begins. + queue.produce("second"); + + queue.shutdown(); + + assertEquals(0, concurrentConsumers.get(), "a consumer was still running after shutdown returned"); + assertEquals(1, maxConcurrentConsumers.get(), "consumer was invoked concurrently during shutdown"); + } + + /** + * The final drain must still flush whatever is left in the partitions once the loops + * have stopped. + */ + @Test + public void testShutdownFinalDrainFlushesRemainder() { + final List consumed = new CopyOnWriteArrayList<>(); + + final BatchQueue queue = BatchQueueManager.create("shutdown-drain-test", + BatchQueueConfig.builder() + .threads(ThreadPolicy.fixed(1)) + .partitions(PartitionPolicy.fixed(1)) + .bufferSize(100) + .minIdleMs(10_000).maxIdleMs(10_000) // keep the drain loop asleep + .shutdownTimeoutMs(2_000) + .consumer(consumed::addAll) + .build()); + + Awaitility.await().atMost(5, TimeUnit.SECONDS) + .until(() -> queue.produce("a") && queue.produce("b")); + + final long start = System.nanoTime(); + queue.shutdown(); + final long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start); + + assertTrue(consumed.containsAll(List.of("a", "b")), "final drain lost data: " + consumed); + // A drain task parked on a long idle backoff must not hold up termination: the scheduler + // drops pending delayed tasks, so shutdown returns well inside its timeout. + assertTrue(elapsedMs < 2_000, + "shutdown waited out its timeout instead of terminating: " + elapsedMs + "ms"); + } + + /** + * A rebalancing queue must still reach termination. On a virtual-thread scheduler the + * periodic task is one submission whose loop exits only on interrupt, so shutdown() has + * to cancel it explicitly or awaitTermination can never succeed. + */ + @Test + public void testShutdownTerminatesWithRebalancingEnabled() { + final BatchQueue queue = BatchQueueManager.create("shutdown-rebalance-test", + BatchQueueConfig.builder() + .threads(ThreadPolicy.fixed(2)) + .partitions(PartitionPolicy.fixed(4)) + .balancer(DrainBalancer.throughputWeighted(), 50) + .bufferSize(100) + .shutdownTimeoutMs(5_000) + .consumer(batch -> { + }) + .build()); + + queue.produce("x"); + final long start = System.nanoTime(); + queue.shutdown(); + final long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start); + + assertTrue(elapsedMs < 5_000, + "shutdown hit the timeout, the periodic rebalance task was not cancelled: " + + elapsedMs + "ms"); + } + + /** + * A timeout must not reintroduce the race. awaitTermination is only a courtesy wait for + * orderly exit; the exclusive dispatch lock is what guarantees the consumer is never entered + * from two threads. Reproduces with a consumer far slower than the timeout. + */ + @Test + public void testTimeoutStillSerialisesTheFinalDispatch() throws Exception { + final CountDownLatch consumerEntered = new CountDownLatch(1); + final AtomicInteger concurrentConsumers = new AtomicInteger(); + final AtomicInteger maxConcurrentConsumers = new AtomicInteger(); + + final BatchQueue queue = BatchQueueManager.create("shutdown-timeout-race-test", + BatchQueueConfig.builder() + .threads(ThreadPolicy.fixed(1)) + .partitions(PartitionPolicy.fixed(1)) + .bufferSize(100) + .shutdownTimeoutMs(50) // far shorter than the consumer below + .consumer(batch -> { + maxConcurrentConsumers.accumulateAndGet( + concurrentConsumers.incrementAndGet(), Math::max); + try { + consumerEntered.countDown(); + Thread.sleep(500); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + concurrentConsumers.decrementAndGet(); + } + }) + .build()); + + queue.produce("first"); + assertTrue(consumerEntered.await(5, TimeUnit.SECONDS), "consumer never started"); + queue.produce("second"); + + queue.shutdown(); + + assertEquals(1, maxConcurrentConsumers.get(), + "final dispatch raced an in-flight consumer after the wait timed out"); + } + + /** + * A caller that loses the shutdown CAS must not return before the winner has finished, or + * shutdownAll() reports completion while a consumer is still running. + */ + @Test + public void testLosingShutdownCallerAwaitsCompletion() throws Exception { + final CountDownLatch consumerEntered = new CountDownLatch(1); + final AtomicInteger activeConsumers = new AtomicInteger(); + + final BatchQueue queue = BatchQueueManager.create("shutdown-share-completion-test", + BatchQueueConfig.builder() + .threads(ThreadPolicy.fixed(1)) + .partitions(PartitionPolicy.fixed(1)) + .bufferSize(100) + .shutdownTimeoutMs(5_000) + .consumer(batch -> { + activeConsumers.incrementAndGet(); + try { + consumerEntered.countDown(); + Thread.sleep(400); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + activeConsumers.decrementAndGet(); + } + }) + .build()); + + queue.produce("first"); + assertTrue(consumerEntered.await(5, TimeUnit.SECONDS), "consumer never started"); + + final Thread winner = new Thread(queue::shutdown); + winner.start(); + Awaitility.await().atMost(2, TimeUnit.SECONDS).until(queue::isShutdownStarted); + + // The loser returns from shutdown(); by then nothing may still be consuming. + queue.shutdown(); + assertEquals(0, activeConsumers.get(), + "the losing shutdown caller returned while a consumer was still running"); + winner.join(10_000); + } + + /** + * onIdle() runs on the drain thread and touches the same worker state as consume() — L1's + * onIdle() calls flush(). It must therefore be inside the dispatch lock, or shutdown's final + * dispatch can run concurrently with it once the courtesy wait expires. + */ + @Test + public void testIdleNotificationIsSerialisedAgainstTheFinalDispatch() throws Exception { + final CountDownLatch idleEntered = new CountDownLatch(1); + final AtomicInteger concurrentCallbacks = new AtomicInteger(); + final AtomicInteger maxConcurrentCallbacks = new AtomicInteger(); + + final HandlerConsumer handler = new HandlerConsumer<>() { + @Override + public void consume(final List data) { + enter(); + try { + Thread.sleep(50); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + concurrentCallbacks.decrementAndGet(); + } + } + + @Override + public void onIdle() { + enter(); + try { + idleEntered.countDown(); + Thread.sleep(500); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + concurrentCallbacks.decrementAndGet(); + } + } + + private void enter() { + maxConcurrentCallbacks.accumulateAndGet(concurrentCallbacks.incrementAndGet(), Math::max); + } + }; + + final BatchQueue queue = BatchQueueManager.create("shutdown-onidle-test", + BatchQueueConfig.builder() + .threads(ThreadPolicy.fixed(1)) + .partitions(PartitionPolicy.fixed(1)) + .bufferSize(100) + .shutdownTimeoutMs(50) // far shorter than the onIdle above + .build()); + queue.addHandler(String.class, handler); + + assertTrue(idleEntered.await(5, TimeUnit.SECONDS), "onIdle never ran"); + queue.produce("queued-behind-idle"); + + queue.shutdown(); + + assertEquals(1, maxConcurrentCallbacks.get(), + "final dispatch ran concurrently with onIdle() after the wait timed out"); + } + + /** + * The winner can exceed shutdownTimeoutMs — it waits on the write lock for as long as the + * consumer runs — so a losing caller must await actual completion, not a fixed bound. + */ + @Test + public void testLosingCallerWaitsOutASlowWinner() throws Exception { + final CountDownLatch consumerEntered = new CountDownLatch(1); + final AtomicInteger activeConsumers = new AtomicInteger(); + + final BatchQueue queue = BatchQueueManager.create("shutdown-slow-winner-test", + BatchQueueConfig.builder() + .threads(ThreadPolicy.fixed(1)) + .partitions(PartitionPolicy.fixed(1)) + .bufferSize(100) + .shutdownTimeoutMs(50) // winner's courtesy wait expires long before... + .consumer(batch -> { + activeConsumers.incrementAndGet(); + try { + consumerEntered.countDown(); + Thread.sleep(1_500); // ...this consumer finishes + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + activeConsumers.decrementAndGet(); + } + }) + .build()); + + queue.produce("first"); + assertTrue(consumerEntered.await(5, TimeUnit.SECONDS), "consumer never started"); + + final Thread winner = new Thread(queue::shutdown); + winner.start(); + Awaitility.await().atMost(2, TimeUnit.SECONDS).until(queue::isShutdownStarted); + + final long loserStart = System.nanoTime(); + queue.shutdown(); // loser + final long loserWaitedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - loserStart); + + assertEquals(0, activeConsumers.get(), + "the losing caller returned while a consumer was still running"); + // The winner's courtesy wait is 50ms but the consumer runs 600ms, so a loser bounded by + // shutdownTimeoutMs returns in ~50ms. Deliberately NOT asserting on winner.isAlive(): + // the latch fires inside doShutdown's finally, so the winner thread is still unwinding + // for a moment afterwards — that is thread teardown, not the contract. + assertTrue(loserWaitedMs >= 100, + "the losing caller did not wait for the winner, returned after " + loserWaitedMs + "ms"); + winner.join(10_000); + } + + /** + * An IO-bound queue behaves identically to a fixed one — the substrate differs, the + * semantics do not. Runs on virtual threads where available, platform threads otherwise. + */ + @Test + public void testIoBoundQueueDrainsAndShutsDown() { + final List consumed = new CopyOnWriteArrayList<>(); + + final BatchQueue queue = BatchQueueManager.create("io-bound-test", + BatchQueueConfig.builder() + .threads(ThreadPolicy.ioBound(4)) + .partitions(PartitionPolicy.fixed(4)) + .bufferSize(10) + .shutdownTimeoutMs(5_000) + .consumer(consumed::addAll) + .build()); + + for (int i = 0; i < 20; i++) { + queue.produce("item-" + i); + } + Awaitility.await().atMost(5, TimeUnit.SECONDS).until(() -> consumed.size() == 20); + + queue.shutdown(); + assertEquals(20, consumed.size()); + } +} diff --git a/oap-server/server-library/library-batch-queue/src/test/java/org/apache/skywalking/oap/server/library/batchqueue/ThreadPolicyTest.java b/oap-server/server-library/library-batch-queue/src/test/java/org/apache/skywalking/oap/server/library/batchqueue/ThreadPolicyTest.java index 3cbdeb6c0aaa..f339bc731dbd 100644 --- a/oap-server/server-library/library-batch-queue/src/test/java/org/apache/skywalking/oap/server/library/batchqueue/ThreadPolicyTest.java +++ b/oap-server/server-library/library-batch-queue/src/test/java/org/apache/skywalking/oap/server/library/batchqueue/ThreadPolicyTest.java @@ -21,6 +21,8 @@ import org.junit.jupiter.api.Test; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -33,6 +35,33 @@ public void testFixedReturnsExactCount() { assertEquals(100, ThreadPolicy.fixed(100).resolve()); } + @Test + public void testIoBoundReturnsExactCount() { + assertEquals(1, ThreadPolicy.ioBound(1).resolve()); + assertEquals(50, ThreadPolicy.ioBound(50).resolve()); + } + + @Test + public void testIoBoundRejectsZero() { + assertThrows(IllegalArgumentException.class, () -> ThreadPolicy.ioBound(0)); + } + + @Test + public void testOnlyIoBoundIsMarkedIoBound() { + assertTrue(ThreadPolicy.ioBound(4).isIoBound()); + assertFalse(ThreadPolicy.fixed(4).isIoBound()); + assertFalse(ThreadPolicy.cpuCores(1.0).isIoBound()); + assertFalse(ThreadPolicy.cpuCoresWithBase(1, 0.25).isIoBound()); + } + + @Test + public void testIoBoundIsNotEqualToFixedWithSameCount() { + // Same resolved count, different substrate intent - they must not collapse. + assertEquals(ThreadPolicy.fixed(4).resolve(), ThreadPolicy.ioBound(4).resolve()); + assertNotEquals(ThreadPolicy.fixed(4), ThreadPolicy.ioBound(4)); + assertEquals("ioBound(4)", ThreadPolicy.ioBound(4).toString()); + } + @Test public void testFixedRejectsZero() { assertThrows(IllegalArgumentException.class, () -> ThreadPolicy.fixed(0));