diff --git a/mateclaw-server/src/main/java/vip/mate/agent/graph/NodeStreamingChatHelper.java b/mateclaw-server/src/main/java/vip/mate/agent/graph/NodeStreamingChatHelper.java index d5ac7042..53b88d5b 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/graph/NodeStreamingChatHelper.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/graph/NodeStreamingChatHelper.java @@ -216,27 +216,6 @@ public class NodeStreamingChatHelper { return m == null ? List.of() : List.of(new vip.mate.llm.failover.FallbackEntry("__legacy__", m)); } - /** - * Record a single primary-model outcome to the health tracker. No-op when - * either the tracker bean isn't wired or the primary's providerId is - * unknown (e.g., tests, legacy callers built without the full constructor). - */ - private void recordPrimary(boolean success) { - if (healthTracker == null || primaryProviderId == null) return; - if (success) healthTracker.recordSuccess(primaryProviderId); - else healthTracker.recordFailure(primaryProviderId); - } - - /** - * Record a primary failure carrying a provider-stated retry window so the - * health tracker can start a cooldown of exactly that length. No-op under - * the same conditions as {@link #recordPrimary}. - */ - private void recordPrimaryFailure(long cooldownOverrideMs) { - if (healthTracker == null || primaryProviderId == null) return; - healthTracker.recordFailure(primaryProviderId, cooldownOverrideMs); - } - /** * Map an {@link ErrorType} to the matching pool * {@link vip.mate.llm.failover.AvailableProviderPool.RemovalSource} for @@ -886,8 +865,8 @@ public class NodeStreamingChatHelper { phase, primaryProviderId, reason); if (broadcast) { broadcastDelta(conversationId, "warning", - buildDeltaJson("主模型暂时不可用（" + (primaryOutOfPool ? "已下线" : "冷却中") - + "），直接尝试备选模型...")); + "主模型暂时不可用（" + (primaryOutOfPool ? "已下线" : "冷却中") + + "），将尝试可用备选模型；如无可用模型请稍后重试。"); } } @@ -923,14 +902,12 @@ public class NodeStreamingChatHelper { } llmCallCount++; if (attempt > 0) retryCount++; - lastResult = doStreamCall(chatModel, prompt, conversationId, phase, publicationMode, attempt, true, retryType, retryHint); + lastResult = doStreamCall(chatModel, prompt, conversationId, phase, publicationMode, attempt, true, retryType, retryHint, primaryProviderId); if (lastResult == null || !lastResult.stopped()) throwIfCancelled(conversationId); if (lastResult != null) { ErrorType errType = lastResult.errorType(); // Success — reaffirm health / pool membership and return. if (lastResult.errorMessage() == null || errType == ErrorType.NONE) { - recordPrimary(true); - addToPool(primaryProviderId); logPerfSummary(phase, conversationId, callStartMs, llmCallCount, retryCount, failoverCount); return lastResult; } @@ -973,22 +950,9 @@ public class NodeStreamingChatHelper { // surface to the caller instead of burning the chain. return lastResult; } - if (errType.countsHealth()) { - // A rate-limit response carrying an explicit retry window - // becomes a health-cooldown override: later turns skip the - // provider until the stated instant instead of re-probing - // it every ~5 minutes and re-collecting the same 429. - Long hintMs = retryHint.get(); - if (errType == ErrorType.RATE_LIMIT && hintMs != null && hintMs > 0) { - recordPrimaryFailure(hintMs); - } else { - recordPrimary(false); - } - healthRecorded = true; - } - if (errType.evictsProvider()) { - removeFromPool(primaryProviderId, errType, lastResult.errorMessage()); - } + // Outcome / pool state is published under the admission permit, + // before queued requests can pass their cooldown check. + healthRecorded = errType.countsHealth(); log.warn("[{}] Primary failed (type={}) — handing off to fallback chain", phase, errType); break; } @@ -1001,7 +965,7 @@ public class NodeStreamingChatHelper { // errors (MODEL_NOT_FOUND et al.) never dent provider health. if (!primarySkipped && !healthRecorded && lastResult != null && isProviderLevelFailure(lastResult.errorType())) { - recordPrimary(false); + if (healthTracker != null) healthTracker.recordFailure(primaryProviderId); } // Primary exhausted retries — walk the fallback chain in priority order. @@ -1034,13 +998,13 @@ public class NodeStreamingChatHelper { fallback.getClass().getSimpleName(), conversationId); if (broadcast) { broadcastDelta(conversationId, "warning", - buildDeltaJson("主模型不可用，正在切换到备选模型 (" + (i + 1) + "/" + fallbackChain.size() + ")...")); + "主模型不可用，正在切换到备选模型 (" + (i + 1) + "/" + fallbackChain.size() + ")..."); } failoverCount++; llmCallCount++; StreamResult fallbackResult = doStreamCall(fallback, prompt, conversationId, phase + "_fallback_" + (i + 1), publicationMode, 0, false, - new AtomicReference<>(), new AtomicReference<>(0L)); + new AtomicReference<>(), new AtomicReference<>(0L), entry.providerId()); if (fallbackResult == null || !fallbackResult.stopped()) throwIfCancelled(conversationId); if (fallbackResult != null && fallbackResult.isPayloadTooLarge()) return fallbackResult; // Accept only fully successful fallbacks. Non-successful results (auth @@ -1049,30 +1013,20 @@ public class NodeStreamingChatHelper { if (fallbackResult != null && fallbackResult.errorType() == ErrorType.NONE && fallbackResult.errorMessage() == null) { - if (healthTracker != null) healthTracker.recordSuccess(entry.providerId()); - addToPool(entry.providerId()); logPerfSummary(phase, conversationId, callStartMs, llmCallCount, retryCount, failoverCount); return fallbackResult; } - // Only provider-level failures count toward the cooldown tracker. A - // null result is a retryable soft failure; a MODEL_NOT_FOUND result is - // model-scoped and must not penalise an otherwise-healthy provider. - if (healthTracker != null - && (fallbackResult == null || isProviderLevelFailure(fallbackResult.errorType()))) { - healthTracker.recordFailure(entry.providerId()); - } if (fallbackResult != null) { - // HARD errors (auth / billing) evict the provider from the pool so - // later walks skip it outright. SOFT and model-scoped errors keep it - // in-pool — absorbed by the tracker's cooldown or simply retried. - removeFromPool(entry.providerId(), fallbackResult.errorType(), fallbackResult.errorMessage()); lastResult = fallbackResult; // remember most recent to report if the whole chain fails } } logPerfSummary(phase, conversationId, callStartMs, llmCallCount, retryCount, failoverCount); + if (lastResult == null && primaryInCooldown && !primaryOutOfPool) + return buildErrorResultWithType("供应商冷却中，且没有可用备选模型，请稍后重试", conversationId, + phase, ErrorType.PROVIDER_BUSY, publicationMode); return lastResult != null ? lastResult - : buildErrorResult("LLM 调用失败，已达最大重试次数", conversationId, phase, publicationMode); + : buildErrorResult("当前没有可用模型，请检查供应商状态或稍后重试", conversationId, phase, publicationMode); } /** D-6: log a structured performance summary for the LLM call phase. */ @@ -1095,7 +1049,7 @@ public class NodeStreamingChatHelper { String conversationId, String phase, PublicationMode publicationMode, int attempt, boolean primaryCall, AtomicReference retryTypeRef, - AtomicReference retryHintRef) { + AtomicReference retryHintRef, String providerId) { // Collapse every SystemMessage in the prompt into a single SystemMessage // at index 0. Some OpenAI-compatible providers (LM Studio's built-in // server, certain strict vLLM / SGLang deployments) reject 400 @@ -1143,8 +1097,43 @@ public class NodeStreamingChatHelper { } } - try { - return doStreamCallInner(chatModel, outbound, conversationId, phase, publicationMode, attempt, primaryCall, retryTypeRef, retryHintRef); + try (var permit = healthTracker == null ? null : healthTracker.acquireRequest(providerId, + () -> streamTracker.isStopRequested(conversationId) || Thread.currentThread().isInterrupted())) { + throwIfCancelled(conversationId); + if (healthTracker != null && permit == null) + return buildErrorResultWithType("供应商请求排队超时，请稍后重试或降低并发", conversationId, + phase, ErrorType.PROVIDER_BUSY, publicationMode); + // Another request may have triggered cooldown while this one queued. + if (!inPool(providerId)) + return buildErrorResultWithType("供应商已下线，请检查配置或切换模型", conversationId, + phase, ErrorType.PROVIDER_BUSY, publicationMode); + if (healthTracker != null && healthTracker.isInCooldown(providerId)) + return buildErrorResultWithType("供应商冷却中，请稍后重试", conversationId, + phase, ErrorType.PROVIDER_BUSY, publicationMode); + StreamResult result = doStreamCallInner(chatModel, outbound, conversationId, phase, + publicationMode, attempt, primaryCall, retryTypeRef, retryHintRef); + // Publish completion state BEFORE releasing capacity. Otherwise a queued + // request can slip between permit release and a newly established cooldown. + if (result != null) { + ErrorType type = result.errorType(); + if (result.errorMessage() == null || type == ErrorType.NONE) { + if (healthTracker != null) healthTracker.recordSuccess(providerId); + addToPool(providerId); + } else { + boolean emptyRetryPending = primaryCall && type == ErrorType.EMPTY_RESPONSE + && attempt < type.retryBudget(); + if (healthTracker != null && type.countsHealth() && !emptyRetryPending) { + Long hint = retryHintRef.get(); + healthTracker.recordFailure(providerId, + type == ErrorType.RATE_LIMIT && hint != null ? hint : 0); + } + removeFromPool(providerId, type, result.errorMessage()); + } + } + return result; + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw cancellation(interrupted); } finally { // Idempotent: if consumer already took the entry, discard is a no-op. if (relayToken != null) { @@ -1213,7 +1202,7 @@ public class NodeStreamingChatHelper { if (broadcast && !emptyResponse) { String cause = overloaded ? "模型服务繁忙" : "请求频率受限"; broadcastDelta(conversationId, "warning", - buildDeltaJson("⏱️ " + cause + "，等待 " + (delay / 1000) + " 秒后重试（第 " + attempt + "/" + MAX_RETRIES + " 次）...")); + "⏱️ " + cause + "，等待 " + (delay / 1000) + " 秒后重试（第 " + attempt + "/" + MAX_RETRIES + " 次）..."); } // Poll stop flag every 100ms so user Stop is honored mid-backoff. long remaining = delay; @@ -1538,7 +1527,7 @@ public class NodeStreamingChatHelper { subscription.dispose(); if (broadcast) { broadcastDelta(conversationId, "warning", - buildDeltaJson("模型在思考阶段停留过久，已自动截断")); + "模型在思考阶段停留过久，已自动截断"); } // dispose 后 latch 可能不会 countDown，直接跳出 break; @@ -1553,7 +1542,7 @@ public class NodeStreamingChatHelper { subscription.dispose(); if (broadcast) { broadcastDelta(conversationId, "warning", - buildDeltaJson("检测到回答内容反复重复，已自动截断")); + "检测到回答内容反复重复，已自动截断"); } break; } @@ -1606,7 +1595,7 @@ public class NodeStreamingChatHelper { phase, contentAccum.length(), toolCallAccumulators.size(), error.getMessage()); if (broadcast) { broadcastDelta(conversationId, "warning", - buildDeltaJson("LLM 响应中断，使用已生成的部分内容继续")); + "LLM 响应中断，使用已生成的部分内容继续"); } return assembleResult(contentAccum, thinkingAccum, toolCallAccumulators, promptTokens.get(), completionTokens.get(), @@ -2033,7 +2022,7 @@ public class NodeStreamingChatHelper { log.error("[{}] Building error result for conversation {}: {}", phase, conversationId, errorMsg); if (publicationMode != PublicationMode.SILENT && streamTracker != null && conversationId != null) { broadcastDelta(conversationId, "warning", - buildDeltaJson(errorMsg)); + errorMsg); } AssistantMessage errorMessage = new AssistantMessage("[错误] " + errorMsg); return new StreamResult("[错误] " + errorMsg, "", errorMessage, @@ -2046,7 +2035,7 @@ public class NodeStreamingChatHelper { log.error("[{}] Building typed error result for conversation {}: {} (type={})", phase, conversationId, errorMsg, errorType); if (publicationMode != PublicationMode.SILENT && streamTracker != null && conversationId != null) { - broadcastDelta(conversationId, "warning", buildDeltaJson(errorMsg)); + broadcastDelta(conversationId, "warning", errorMsg); // 广播结构化 error 事件，供前端展示错误卡片 String errorJson = buildErrorEventJson(errorMsg, conversationId, errorType); streamTracker.broadcast(conversationId, "error", errorJson); @@ -2236,6 +2225,8 @@ public class NodeStreamingChatHelper { * staying on a rate-limited provider wastes time — then fail over. */ RATE_LIMIT (MAX_RETRIES_RATE_LIMIT, true, false, true), + /** Local admission/cooldown; no provider request failed and health must not be penalized. */ + PROVIDER_BUSY (0, true, false, false), /** * The provider's serving capacity is saturated (HTTP 529, * "engine_overloaded", "model is overloaded"). The caller's key is diff --git a/mateclaw-server/src/main/java/vip/mate/llm/failover/ProviderHealthProperties.java b/mateclaw-server/src/main/java/vip/mate/llm/failover/ProviderHealthProperties.java index 5b4ea1b3..a034afc6 100644 --- a/mateclaw-server/src/main/java/vip/mate/llm/failover/ProviderHealthProperties.java +++ b/mateclaw-server/src/main/java/vip/mate/llm/failover/ProviderHealthProperties.java @@ -29,6 +29,17 @@ public class ProviderHealthProperties { /** Cooldown window in milliseconds. */ private long cooldownMs = 300_000L; + /** Shared per-provider limit for agent LLM streams, including fallbacks. */ + private int maxConcurrentRequests = 1; + + /** Maximum admission wait; queued callers remain cancellable. */ + private long admissionTimeoutMs = 120_000L; + + public int getMaxConcurrentRequests() { return maxConcurrentRequests; } + public void setMaxConcurrentRequests(int value) { maxConcurrentRequests = Math.max(1, value); } + public long getAdmissionTimeoutMs() { return admissionTimeoutMs; } + public void setAdmissionTimeoutMs(long value) { admissionTimeoutMs = Math.max(1, value); } + /** * TTL after which a provider HARD-removed for BILLING is lazily readmitted * to the available pool. Users top up balances and aggregator quotas diff --git a/mateclaw-server/src/main/java/vip/mate/llm/failover/ProviderHealthTracker.java b/mateclaw-server/src/main/java/vip/mate/llm/failover/ProviderHealthTracker.java index eb1a5e2a..4bbfd6e8 100644 --- a/mateclaw-server/src/main/java/vip/mate/llm/failover/ProviderHealthTracker.java +++ b/mateclaw-server/src/main/java/vip/mate/llm/failover/ProviderHealthTracker.java @@ -39,6 +39,39 @@ public class ProviderHealthTracker { private final ProviderHealthProperties props; private final ConcurrentHashMap consecutiveFailures = new ConcurrentHashMap<>(); private final ConcurrentHashMap cooldownUntilMs = new ConcurrentHashMap<>(); + private final ConcurrentHashMap requestSlots = new ConcurrentHashMap<>(); + + /** Acquire across helper instances; hold until the entire stream terminates, not just headers. */ + public RequestPermit acquireRequest(String providerId, java.util.function.BooleanSupplier cancelled) + throws InterruptedException { + if (providerId == null || providerId.isBlank()) return new RequestPermit(null); + var slots = requestSlots.computeIfAbsent(providerId, + ignored -> new java.util.concurrent.Semaphore(props.getMaxConcurrentRequests(), true)); + long deadline = System.nanoTime() + java.util.concurrent.TimeUnit.MILLISECONDS.toNanos(props.getAdmissionTimeoutMs()); + while (true) { + if (cancelled.getAsBoolean()) throw new java.util.concurrent.CancellationException("Queued LLM call cancelled"); + long remaining = deadline - System.nanoTime(); + if (remaining <= 0) return null; + if (slots.tryAcquire(Math.min(remaining, java.util.concurrent.TimeUnit.MILLISECONDS.toNanos(100)), + java.util.concurrent.TimeUnit.NANOSECONDS)) { + var permit = new RequestPermit(slots); + if (cancelled.getAsBoolean()) { + permit.close(); + throw new java.util.concurrent.CancellationException("Queued LLM call cancelled"); + } + return permit; + } + } + } + + public static final class RequestPermit implements AutoCloseable { + private final java.util.concurrent.Semaphore slots; + private final java.util.concurrent.atomic.AtomicBoolean closed = new java.util.concurrent.atomic.AtomicBoolean(); + private RequestPermit(java.util.concurrent.Semaphore slots) { this.slots = slots; } + @Override public void close() { + if (slots != null && closed.compareAndSet(false, true)) slots.release(); + } + } public ProviderHealthTracker(ProviderHealthProperties props) { this.props = props; diff --git a/mateclaw-server/src/main/resources/application.yml b/mateclaw-server/src/main/resources/application.yml index 88eecc6a..d8bc30dc 100644 --- a/mateclaw-server/src/main/resources/application.yml +++ b/mateclaw-server/src/main/resources/application.yml @@ -404,6 +404,9 @@ mateclaw: enabled: true failure-threshold: 3 cooldown-ms: 300000 + # Shared by agent conversations and fallback models of the same provider. + max-concurrent-requests: ${MATECLAW_LLM_MAX_CONCURRENT_REQUESTS:1} + admission-timeout-ms: ${MATECLAW_LLM_ADMISSION_TIMEOUT_MS:120000} # Multi-agent delegation (DelegateAgentTool). delegation: # Wall-clock budget for one delegateParallel batch (shared across all diff --git a/mateclaw-server/src/test/java/vip/mate/agent/graph/NodeStreamingChatHelperAdmissionTest.java b/mateclaw-server/src/test/java/vip/mate/agent/graph/NodeStreamingChatHelperAdmissionTest.java new file mode 100644 index 00000000..f98b9ddd --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/agent/graph/NodeStreamingChatHelperAdmissionTest.java @@ -0,0 +1,121 @@ +package vip.mate.agent.graph; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.model.*; +import org.springframework.ai.chat.prompt.Prompt; +import reactor.core.publisher.Flux; +import vip.mate.channel.web.ChatStreamTracker; +import vip.mate.llm.failover.*; + +import java.time.Duration; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; +import static org.mockito.ArgumentMatchers.*; + +class NodeStreamingChatHelperAdmissionTest { + @Test + void cooldownIsPublishedBeforeNextQueuedRequestCanStart() throws Exception { + var recording = new CountDownLatch(1); + var release = new CountDownLatch(1); + var props = new ProviderHealthProperties(); + props.setFailureThreshold(1); + var health = new ProviderHealthTracker(props) { + @Override public void recordFailure(String id, long hint) { + recording.countDown(); + try { assertTrue(release.await(3, TimeUnit.SECONDS)); } + catch (InterruptedException e) { throw new AssertionError(e); } + super.recordFailure(id, hint); + } + }; + var tracker = mock(ChatStreamTracker.class); + var bad = mock(ChatModel.class); + when(bad.stream(any(Prompt.class))).thenReturn(Flux.error(new RuntimeException("401 Unauthorized"))); + var unexpectedCall = new CountDownLatch(1); + var next = mock(ChatModel.class); + when(next.stream(any(Prompt.class))).thenAnswer(ignored -> { + unexpectedCall.countDown(); + return Flux.just(new ChatResponse(List.of(new Generation(new AssistantMessage("unexpected"))))); + }); + try (var workers = Executors.newFixedThreadPool(2)) { + var first = workers.submit(() -> new NodeStreamingChatHelper(tracker, List.of(), null, health, "same") + .streamCall(bad, new Prompt("hi"), "first", "reasoning")); + assertTrue(recording.await(1, TimeUnit.SECONDS)); + var second = workers.submit(() -> new NodeStreamingChatHelper(tracker, List.of(), null, health, "same") + .streamCall(next, new Prompt("hi"), "next", "reasoning")); + try { assertFalse(unexpectedCall.await(150, TimeUnit.MILLISECONDS)); } + finally { release.countDown(); } + first.get(3, TimeUnit.SECONDS); + assertEquals(NodeStreamingChatHelper.ErrorType.PROVIDER_BUSY, + second.get(3, TimeUnit.SECONDS).errorType()); + } finally { release.countDown(); } + verify(next, never()).stream(any(Prompt.class)); + } + + @Test + void separateHelpersShareProviderAdmissionForWholeStream() throws Exception { + exerciseSharedAdmission(false); + } + + @Test + void fallbackStreamsUseTheSameSharedAdmission() throws Exception { + exerciseSharedAdmission(true); + } + + private void exerciseSharedAdmission(boolean fallback) throws Exception { + var tracker = mock(ChatStreamTracker.class); + var health = new ProviderHealthTracker(new ProviderHealthProperties()); + var active = new AtomicInteger(); + var peak = new AtomicInteger(); + var model = mock(ChatModel.class); + var response = new ChatResponse(List.of(new Generation(new AssistantMessage("ok")))); + when(model.stream(any(Prompt.class))).thenAnswer(ignored -> Flux.defer(() -> { + peak.accumulateAndGet(active.incrementAndGet(), Math::max); + return Flux.just(response).delayElements(Duration.ofMillis(250)) + .doOnComplete(active::decrementAndGet); + })); + var start = new CountDownLatch(1); + try (var workers = Executors.newFixedThreadPool(2)) { + var futures = new java.util.ArrayList>(); + for (int i = 0; i < 2; i++) { + String id = "admission-" + i; + var primary = fallback ? mock(ChatModel.class) : model; + if (fallback) when(primary.stream(any(Prompt.class))) + .thenReturn(Flux.error(new RuntimeException("401 Unauthorized"))); + var helper = new NodeStreamingChatHelper(tracker, + fallback ? List.of(new FallbackEntry("shared", model)) : List.of(), + null, health, fallback ? "primary-" + i : "shared"); + futures.add(workers.submit(() -> { + start.await(); + assertEquals("ok", helper.streamCall(primary, new Prompt("hello"), id, "reasoning").text()); + return null; + })); + } + start.countDown(); + for (var future : futures) future.get(5, TimeUnit.SECONDS); + } + assertEquals(1, peak.get(), "Same-provider streams must not overlap across helpers"); + } + + @Test + void cooldownWithoutFallbackIsTypedAndWarningHasOneJsonEnvelope() throws Exception { + var tracker = mock(ChatStreamTracker.class); + var health = new ProviderHealthTracker(new ProviderHealthProperties()); + health.recordFailure("limited", 60_000); + var model = mock(ChatModel.class); + var helper = new NodeStreamingChatHelper(tracker, List.of(), null, health, "limited"); + var result = helper.streamCall(model, new Prompt("hello"), "cooldown", "plan_summary"); + assertEquals(NodeStreamingChatHelper.ErrorType.PROVIDER_BUSY, result.errorType()); + assertFalse(result.errorMessage().contains("最大重试"), "No request was attempted"); + verify(model, never()).stream(any(Prompt.class)); + var payloads = org.mockito.ArgumentCaptor.forClass(String.class); + verify(tracker, atLeastOnce()).broadcast(eq("cooldown"), eq("warning"), payloads.capture()); + var json = new com.fasterxml.jackson.databind.ObjectMapper(); + for (String payload : payloads.getAllValues()) + assertFalse(json.readTree(payload).path("delta").asText().startsWith("{")); + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/llm/failover/ProviderRequestAdmissionTest.java b/mateclaw-server/src/test/java/vip/mate/llm/failover/ProviderRequestAdmissionTest.java new file mode 100644 index 00000000..e008ce4a --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/llm/failover/ProviderRequestAdmissionTest.java @@ -0,0 +1,46 @@ +package vip.mate.llm.failover; + +import org.junit.jupiter.api.Test; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicBoolean; +import static org.junit.jupiter.api.Assertions.*; + +class ProviderRequestAdmissionTest { + @Test + void differentProvidersProceedAndQueueTimeoutDoesNotDamageHealth() throws Exception { + var props = new ProviderHealthProperties(); + props.setAdmissionTimeoutMs(20); + var tracker = new ProviderHealthTracker(props); + try (var first = tracker.acquireRequest("a", () -> false); + var independent = tracker.acquireRequest("b", () -> false)) { + assertNotNull(first); + assertNotNull(independent); + assertNull(tracker.acquireRequest("a", () -> false)); + assertTrue(tracker.snapshot().isEmpty()); + } + try (var reused = tracker.acquireRequest("a", () -> false)) { + assertNotNull(reused, "Closing the first stream must release its permit"); + } + } + + @Test + void queuedCancellationDoesNotStealOrLeakCapacity() throws Exception { + var tracker = new ProviderHealthTracker(new ProviderHealthProperties()); + var cancel = new AtomicBoolean(); + try (var workers = Executors.newSingleThreadExecutor(); + var first = tracker.acquireRequest("a", () -> false)) { + var entered = new CountDownLatch(1); + var future = workers.submit(() -> { + entered.countDown(); + return tracker.acquireRequest("a", cancel::get); + }); + assertTrue(entered.await(1, TimeUnit.SECONDS)); + cancel.set(true); + var failure = assertThrows(ExecutionException.class, () -> future.get(1, TimeUnit.SECONDS)); + assertInstanceOf(CancellationException.class, failure.getCause()); + } + try (var reused = tracker.acquireRequest("a", () -> false)) { + assertNotNull(reused); + } + } +}