mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-05 23:02:10 +03:00
Merge pull request #1524 from kang-heewon/provider-account-concurrency-cap
feat: provider/account-level concurrency cap enforcement
This commit is contained in:
Diego Rodrigues de Sa e Souza
@@ -2347,6 +2347,10 @@ components:
|
||||
type: string
|
||||
isActive:
|
||||
type: boolean
|
||||
maxConcurrent:
|
||||
type: integer
|
||||
nullable: true
|
||||
minimum: 0
|
||||
priority:
|
||||
type: integer
|
||||
testStatus:
|
||||
@@ -2372,6 +2376,10 @@ components:
|
||||
isActive:
|
||||
type: boolean
|
||||
default: true
|
||||
maxConcurrent:
|
||||
type: integer
|
||||
nullable: true
|
||||
minimum: 0
|
||||
|
||||
ApiKey:
|
||||
type: object
|
||||
|
||||
@@ -90,6 +90,11 @@ import {
|
||||
updateFromResponseBody,
|
||||
initializeRateLimits,
|
||||
} from "../services/rateLimitManager.ts";
|
||||
import {
|
||||
acquire as acquireAccountSemaphore,
|
||||
buildAccountSemaphoreKey,
|
||||
markBlocked as markAccountSemaphoreBlocked,
|
||||
} from "../services/accountSemaphore.ts";
|
||||
import {
|
||||
generateSignature,
|
||||
getCachedResponse,
|
||||
@@ -465,6 +470,72 @@ function getHeaderValueCaseInsensitive(
|
||||
return null;
|
||||
}
|
||||
|
||||
function toFiniteNumberOrNull(value: unknown): number | null {
|
||||
if (typeof value === "number" && Number.isFinite(value)) {
|
||||
return value;
|
||||
}
|
||||
if (typeof value === "string" && value.trim().length > 0) {
|
||||
const parsed = Number(value);
|
||||
return Number.isFinite(parsed) ? parsed : null;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
function isSemaphoreTimeoutError(error: unknown): error is Error & { code: string } {
|
||||
return (
|
||||
!!error &&
|
||||
typeof error === "object" &&
|
||||
(error as { code?: unknown }).code === "SEMAPHORE_TIMEOUT"
|
||||
);
|
||||
}
|
||||
|
||||
function resolveAccountSemaphoreAccountKey(
|
||||
connectionId: string | null | undefined,
|
||||
credentials: Record<string, unknown> | null | undefined
|
||||
): string | null {
|
||||
if (typeof connectionId === "string" && connectionId.trim().length > 0) {
|
||||
return connectionId;
|
||||
}
|
||||
|
||||
const candidateKeys = [
|
||||
credentials?.connectionId,
|
||||
credentials?.id,
|
||||
credentials?.email,
|
||||
credentials?.name,
|
||||
credentials?.displayName,
|
||||
];
|
||||
|
||||
for (const candidate of candidateKeys) {
|
||||
if (typeof candidate === "string" && candidate.trim().length > 0) {
|
||||
return candidate.trim();
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
function resolveAccountSemaphoreMaxConcurrency(
|
||||
credentials: Record<string, unknown> | null | undefined
|
||||
): number | null {
|
||||
return toFiniteNumberOrNull(credentials?.maxConcurrent);
|
||||
}
|
||||
|
||||
function resolveAccountSemaphoreKey({
|
||||
provider,
|
||||
model,
|
||||
connectionId,
|
||||
credentials,
|
||||
}: {
|
||||
provider: string | null | undefined;
|
||||
model: string;
|
||||
connectionId: string | null | undefined;
|
||||
credentials: Record<string, unknown> | null | undefined;
|
||||
}): string | null {
|
||||
const accountKey = resolveAccountSemaphoreAccountKey(connectionId, credentials);
|
||||
if (!accountKey || !provider) return null;
|
||||
return buildAccountSemaphoreKey({ provider, accountKey });
|
||||
}
|
||||
|
||||
function buildClaudePromptCacheLogMeta(
|
||||
targetFormat: string,
|
||||
finalBody: Record<string, unknown> | null | undefined,
|
||||
@@ -1783,6 +1854,15 @@ export async function handleChatCore({
|
||||
|
||||
const executeProviderRequest = async (modelToCall = effectiveModel, allowDedup = false) => {
|
||||
const execute = async () => {
|
||||
const executionCredentials = getExecutionCredentials();
|
||||
const accountSemaphoreMaxConcurrency =
|
||||
resolveAccountSemaphoreMaxConcurrency(executionCredentials);
|
||||
const accountSemaphoreKey = resolveAccountSemaphoreKey({
|
||||
provider,
|
||||
model: modelToCall,
|
||||
connectionId,
|
||||
credentials: executionCredentials,
|
||||
});
|
||||
let bodyToSend =
|
||||
translatedBody.model === modelToCall
|
||||
? translatedBody
|
||||
@@ -1850,6 +1930,13 @@ export async function handleChatCore({
|
||||
}
|
||||
}
|
||||
|
||||
const acquireAccountSemaphoreRelease =
|
||||
accountSemaphoreKey && accountSemaphoreMaxConcurrency != null
|
||||
? await acquireAccountSemaphore(accountSemaphoreKey, {
|
||||
maxConcurrency: accountSemaphoreMaxConcurrency,
|
||||
})
|
||||
: () => {};
|
||||
|
||||
const rawResult = await withRateLimit(provider, connectionId, modelToCall, async () => {
|
||||
let attempts = 0;
|
||||
const maxAttempts = provider === "qwen" ? 3 : 1;
|
||||
@@ -1859,7 +1946,7 @@ export async function handleChatCore({
|
||||
model: modelToCall,
|
||||
body: bodyToSend,
|
||||
stream: upstreamStream,
|
||||
credentials: getExecutionCredentials(),
|
||||
credentials: executionCredentials,
|
||||
signal: streamController.signal,
|
||||
log,
|
||||
extendedContext,
|
||||
@@ -1882,17 +1969,39 @@ export async function handleChatCore({
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// For streaming: wrap response body to release semaphore only when stream is fully consumed
|
||||
if (stream) {
|
||||
const originalBody = res.response.body;
|
||||
const wrappedBody = originalBody
|
||||
? originalBody.pipeThrough(
|
||||
new TransformStream({
|
||||
flush: () => {
|
||||
acquireAccountSemaphoreRelease();
|
||||
},
|
||||
})
|
||||
)
|
||||
: null;
|
||||
return {
|
||||
...res,
|
||||
response: new Response(wrappedBody, {
|
||||
status: res.response.status,
|
||||
statusText: res.response.statusText,
|
||||
headers: res.response.headers,
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
});
|
||||
|
||||
if (stream) return rawResult;
|
||||
|
||||
// Non-stream responses need cloning for shared dedup consumers.
|
||||
// Non-stream: release semaphore immediately after reading full response body
|
||||
const status = rawResult.response.status;
|
||||
const statusText = rawResult.response.statusText;
|
||||
const headers = Array.from(rawResult.response.headers.entries()) as [string, string][];
|
||||
const payload = await rawResult.response.text();
|
||||
acquireAccountSemaphoreRelease();
|
||||
|
||||
return {
|
||||
...rawResult,
|
||||
@@ -1967,6 +2076,28 @@ export async function handleChatCore({
|
||||
);
|
||||
} catch (error) {
|
||||
trackPendingRequest(model, provider, connectionId, false);
|
||||
if (isSemaphoreTimeoutError(error)) {
|
||||
appendRequestLog({
|
||||
model,
|
||||
provider,
|
||||
connectionId,
|
||||
status: `FAILED ${error.code}`,
|
||||
}).catch(() => {});
|
||||
if (isCombo) {
|
||||
throw error;
|
||||
}
|
||||
const failureMessage = error.message || "Semaphore timeout";
|
||||
persistAttemptLogs({
|
||||
status: HTTP_STATUS.RATE_LIMITED,
|
||||
error: failureMessage,
|
||||
providerRequest: finalBody || translatedBody,
|
||||
clientResponse: buildErrorBody(HTTP_STATUS.RATE_LIMITED, failureMessage),
|
||||
claudeCacheMeta: claudePromptCacheLogMeta,
|
||||
cacheSource: "upstream",
|
||||
});
|
||||
persistFailureUsage(HTTP_STATUS.RATE_LIMITED, error.code);
|
||||
return createErrorResult(HTTP_STATUS.RATE_LIMITED, failureMessage);
|
||||
}
|
||||
const failureStatus =
|
||||
error.name === "AbortError"
|
||||
? 499
|
||||
@@ -2138,6 +2269,80 @@ export async function handleChatCore({
|
||||
console.warn(
|
||||
`[provider] Node ${connectionId} account deactivated (${statusCode}) — disabling permanently`
|
||||
);
|
||||
} else if (errorType === PROVIDER_ERROR_TYPES.RATE_LIMITED) {
|
||||
// For providers with per-model quotas (passthrough providers, Gemini),
|
||||
// each model has independent quota. A 429 on one model must NOT lock out
|
||||
// the entire connection — other models may still have quota available.
|
||||
const rateLimitCooldownMs = retryAfterMs || COOLDOWN_MS.rateLimit;
|
||||
const accountSemaphoreKey = resolveAccountSemaphoreKey({
|
||||
provider,
|
||||
model: currentModel,
|
||||
connectionId,
|
||||
credentials,
|
||||
});
|
||||
if (accountSemaphoreKey) {
|
||||
markAccountSemaphoreBlocked(accountSemaphoreKey, rateLimitCooldownMs);
|
||||
}
|
||||
if (
|
||||
lockModelIfPerModelQuota(
|
||||
provider,
|
||||
connectionId,
|
||||
model,
|
||||
"rate_limited",
|
||||
rateLimitCooldownMs
|
||||
)
|
||||
) {
|
||||
console.warn(
|
||||
`[provider] Node ${connectionId} model-only rate limited (${statusCode}) for ${model} - ${Math.ceil(rateLimitCooldownMs / 1000)}s (connection stays active)`
|
||||
);
|
||||
} else {
|
||||
const rateLimitedUntil = new Date(Date.now() + rateLimitCooldownMs).toISOString();
|
||||
await updateProviderConnection(connectionId, {
|
||||
rateLimitedUntil: rateLimitedUntil,
|
||||
testStatus: "unavailable",
|
||||
lastErrorType: errorType,
|
||||
lastError: message,
|
||||
errorCode: statusCode,
|
||||
healthCheckInterval: null,
|
||||
lastHealthCheckAt: null,
|
||||
});
|
||||
console.warn(
|
||||
`[provider] Node ${connectionId} rate limited (${statusCode}) - Next available at ${rateLimitedUntil}`
|
||||
);
|
||||
}
|
||||
} else if (errorType === PROVIDER_ERROR_TYPES.QUOTA_EXHAUSTED) {
|
||||
// Providers with per-model quotas — lock the model only, not the connection
|
||||
const quotaCooldownMs = retryAfterMs || COOLDOWN_MS.rateLimit;
|
||||
const accountSemaphoreKey = resolveAccountSemaphoreKey({
|
||||
provider,
|
||||
model: currentModel,
|
||||
connectionId,
|
||||
credentials,
|
||||
});
|
||||
if (accountSemaphoreKey) {
|
||||
markAccountSemaphoreBlocked(accountSemaphoreKey, quotaCooldownMs);
|
||||
}
|
||||
if (
|
||||
lockModelIfPerModelQuota(
|
||||
provider,
|
||||
connectionId,
|
||||
model,
|
||||
"quota_exhausted",
|
||||
quotaCooldownMs
|
||||
)
|
||||
) {
|
||||
console.warn(
|
||||
`[provider] Node ${connectionId} model-only quota exhausted (${statusCode}) for ${model} - ${Math.ceil(quotaCooldownMs / 1000)}s (connection stays active)`
|
||||
);
|
||||
} else {
|
||||
await updateProviderConnection(connectionId, {
|
||||
testStatus: "credits_exhausted",
|
||||
lastErrorType: errorType,
|
||||
lastError: message,
|
||||
errorCode: statusCode,
|
||||
});
|
||||
console.warn(`[provider] Node ${connectionId} exhausted quota (${statusCode})`);
|
||||
}
|
||||
} else if (errorType === PROVIDER_ERROR_TYPES.ACCOUNT_DEACTIVATED) {
|
||||
await updateProviderConnection(connectionId, {
|
||||
isActive: false,
|
||||
|
||||
288
open-sse/services/accountSemaphore.ts
Normal file
288
open-sse/services/accountSemaphore.ts
Normal file
@@ -0,0 +1,288 @@
|
||||
/**
|
||||
* Account Semaphore
|
||||
*
|
||||
* In-memory provider/account concurrency limiter keyed by provider and account.
|
||||
* Requests beyond the configured concurrency cap wait in a FIFO queue until a slot opens,
|
||||
* the gate is unblocked, or the queue timeout expires.
|
||||
*/
|
||||
|
||||
export interface AccountSemaphoreKeyParts {
|
||||
provider: string;
|
||||
accountKey: string;
|
||||
}
|
||||
|
||||
interface QueuedAcquire {
|
||||
resolve: (release: () => void) => void;
|
||||
reject: (error: Error) => void;
|
||||
timer: ReturnType<typeof setTimeout>;
|
||||
}
|
||||
|
||||
interface AccountGate {
|
||||
running: number;
|
||||
maxConcurrency: number;
|
||||
queue: QueuedAcquire[];
|
||||
blockedUntil: number | null;
|
||||
cleanupTimer: ReturnType<typeof setTimeout> | null;
|
||||
}
|
||||
|
||||
export interface AcquireAccountSemaphoreOptions {
|
||||
maxConcurrency?: number | null;
|
||||
timeoutMs?: number;
|
||||
}
|
||||
|
||||
export interface AccountSemaphoreStatsEntry {
|
||||
running: number;
|
||||
queued: number;
|
||||
maxConcurrency: number;
|
||||
blockedUntil: string | null;
|
||||
}
|
||||
|
||||
const DEFAULT_TIMEOUT_MS = 30_000;
|
||||
|
||||
const gates = new Map<string, AccountGate>();
|
||||
|
||||
/**
|
||||
* Build the canonical account semaphore key.
|
||||
*/
|
||||
export function buildAccountSemaphoreKey({
|
||||
provider,
|
||||
accountKey,
|
||||
}: AccountSemaphoreKeyParts): string {
|
||||
return `${String(provider)}:${String(accountKey)}`;
|
||||
}
|
||||
|
||||
function isBypassed(maxConcurrency?: number | null): boolean {
|
||||
return maxConcurrency == null || maxConcurrency <= 0;
|
||||
}
|
||||
|
||||
function createNoopReleaseFn(): () => void {
|
||||
let released = false;
|
||||
|
||||
return () => {
|
||||
if (released) return;
|
||||
released = true;
|
||||
};
|
||||
}
|
||||
|
||||
function ensureGate(semaphoreKey: string, maxConcurrency: number): AccountGate {
|
||||
const existing = gates.get(semaphoreKey);
|
||||
if (existing) {
|
||||
existing.maxConcurrency = maxConcurrency;
|
||||
return existing;
|
||||
}
|
||||
|
||||
const created: AccountGate = {
|
||||
running: 0,
|
||||
maxConcurrency,
|
||||
queue: [],
|
||||
blockedUntil: null,
|
||||
cleanupTimer: null,
|
||||
};
|
||||
gates.set(semaphoreKey, created);
|
||||
return created;
|
||||
}
|
||||
|
||||
function isBlocked(gate: AccountGate): boolean {
|
||||
if (!gate.blockedUntil) return false;
|
||||
if (Date.now() >= gate.blockedUntil) {
|
||||
gate.blockedUntil = null;
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
function clearCleanupTimer(gate: AccountGate): void {
|
||||
if (!gate.cleanupTimer) return;
|
||||
clearTimeout(gate.cleanupTimer);
|
||||
gate.cleanupTimer = null;
|
||||
}
|
||||
|
||||
function cleanupGateIfIdle(semaphoreKey: string): void {
|
||||
const gate = gates.get(semaphoreKey);
|
||||
if (!gate) return;
|
||||
if (gate.running > 0 || gate.queue.length > 0 || isBlocked(gate)) return;
|
||||
clearCleanupTimer(gate);
|
||||
gates.delete(semaphoreKey);
|
||||
}
|
||||
|
||||
function scheduleCleanup(semaphoreKey: string): void {
|
||||
const gate = gates.get(semaphoreKey);
|
||||
if (!gate) return;
|
||||
clearCleanupTimer(gate);
|
||||
|
||||
gate.cleanupTimer = setTimeout(() => {
|
||||
gate.cleanupTimer = null;
|
||||
cleanupGateIfIdle(semaphoreKey);
|
||||
}, 0);
|
||||
|
||||
gate.cleanupTimer.unref?.();
|
||||
}
|
||||
|
||||
function drainQueue(semaphoreKey: string): void {
|
||||
const gate = gates.get(semaphoreKey);
|
||||
if (!gate) return;
|
||||
|
||||
while (gate.queue.length > 0 && gate.running < gate.maxConcurrency && !isBlocked(gate)) {
|
||||
const next = gate.queue.shift();
|
||||
if (!next) break;
|
||||
clearTimeout(next.timer);
|
||||
gate.running++;
|
||||
next.resolve(createReleaseFn(semaphoreKey));
|
||||
}
|
||||
|
||||
if (gate.running === 0 && gate.queue.length === 0) {
|
||||
scheduleCleanup(semaphoreKey);
|
||||
}
|
||||
}
|
||||
|
||||
function createReleaseFn(semaphoreKey: string): () => void {
|
||||
let released = false;
|
||||
|
||||
return () => {
|
||||
if (released) return;
|
||||
released = true;
|
||||
|
||||
const gate = gates.get(semaphoreKey);
|
||||
if (!gate) return;
|
||||
if (gate.running > 0) {
|
||||
gate.running--;
|
||||
}
|
||||
|
||||
if (gate.queue.length > 0) {
|
||||
drainQueue(semaphoreKey);
|
||||
return;
|
||||
}
|
||||
|
||||
scheduleCleanup(semaphoreKey);
|
||||
};
|
||||
}
|
||||
|
||||
function createSemaphoreTimeoutError(
|
||||
semaphoreKey: string,
|
||||
timeoutMs: number
|
||||
): Error & { code: string } {
|
||||
const error = new Error(`Semaphore timeout after ${timeoutMs}ms for ${semaphoreKey}`) as Error & {
|
||||
code: string;
|
||||
};
|
||||
error.code = "SEMAPHORE_TIMEOUT";
|
||||
return error;
|
||||
}
|
||||
|
||||
/**
|
||||
* Acquire a slot for a provider/model/account tuple.
|
||||
* Returns an idempotent release function that is safe to call in finally blocks.
|
||||
*/
|
||||
export function acquire(
|
||||
semaphoreKey: string,
|
||||
{ maxConcurrency = null, timeoutMs = DEFAULT_TIMEOUT_MS }: AcquireAccountSemaphoreOptions = {}
|
||||
): Promise<() => void> {
|
||||
if (isBypassed(maxConcurrency)) {
|
||||
return Promise.resolve(createNoopReleaseFn());
|
||||
}
|
||||
|
||||
const gate = ensureGate(semaphoreKey, maxConcurrency);
|
||||
clearCleanupTimer(gate);
|
||||
|
||||
if (gate.running < gate.maxConcurrency && !isBlocked(gate)) {
|
||||
gate.running++;
|
||||
return Promise.resolve(createReleaseFn(semaphoreKey));
|
||||
}
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
const timer = setTimeout(() => {
|
||||
const nextGate = gates.get(semaphoreKey);
|
||||
if (!nextGate) {
|
||||
reject(createSemaphoreTimeoutError(semaphoreKey, timeoutMs));
|
||||
return;
|
||||
}
|
||||
|
||||
const queueIndex = nextGate.queue.findIndex((item) => item.timer === timer);
|
||||
if (queueIndex !== -1) {
|
||||
nextGate.queue.splice(queueIndex, 1);
|
||||
}
|
||||
|
||||
if (nextGate.running === 0 && nextGate.queue.length === 0) {
|
||||
scheduleCleanup(semaphoreKey);
|
||||
}
|
||||
|
||||
reject(createSemaphoreTimeoutError(semaphoreKey, timeoutMs));
|
||||
}, timeoutMs);
|
||||
|
||||
timer.unref?.();
|
||||
gate.queue.push({ resolve, reject, timer });
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Temporarily block new acquisitions for a key while allowing in-flight requests to finish.
|
||||
*/
|
||||
export function markBlocked(semaphoreKey: string, cooldownMs: number): void {
|
||||
const safeCooldownMs = Number.isFinite(cooldownMs) && cooldownMs > 0 ? cooldownMs : 0;
|
||||
if (safeCooldownMs <= 0) {
|
||||
const gate = gates.get(semaphoreKey);
|
||||
if (!gate) return;
|
||||
gate.blockedUntil = null;
|
||||
drainQueue(semaphoreKey);
|
||||
return;
|
||||
}
|
||||
|
||||
const gate = gates.get(semaphoreKey) ?? ensureGate(semaphoreKey, 1);
|
||||
clearCleanupTimer(gate);
|
||||
gate.blockedUntil = Date.now() + safeCooldownMs;
|
||||
|
||||
const timer = setTimeout(() => {
|
||||
const nextGate = gates.get(semaphoreKey);
|
||||
if (!nextGate) return;
|
||||
if (nextGate.blockedUntil && Date.now() >= nextGate.blockedUntil) {
|
||||
nextGate.blockedUntil = null;
|
||||
drainQueue(semaphoreKey);
|
||||
if (nextGate.running === 0 && nextGate.queue.length === 0) {
|
||||
scheduleCleanup(semaphoreKey);
|
||||
}
|
||||
}
|
||||
}, safeCooldownMs + 50);
|
||||
|
||||
timer.unref?.();
|
||||
}
|
||||
|
||||
/**
|
||||
* Return the current in-memory semaphore snapshot.
|
||||
*/
|
||||
export function getStats(): Record<string, AccountSemaphoreStatsEntry> {
|
||||
const stats: Record<string, AccountSemaphoreStatsEntry> = {};
|
||||
|
||||
for (const [key, gate] of gates) {
|
||||
stats[key] = {
|
||||
running: gate.running,
|
||||
queued: gate.queue.length,
|
||||
maxConcurrency: gate.maxConcurrency,
|
||||
blockedUntil: gate.blockedUntil ? new Date(gate.blockedUntil).toISOString() : null,
|
||||
};
|
||||
}
|
||||
|
||||
return stats;
|
||||
}
|
||||
|
||||
/**
|
||||
* Reset a single key and reject queued waiters.
|
||||
*/
|
||||
export function reset(semaphoreKey: string): void {
|
||||
const gate = gates.get(semaphoreKey);
|
||||
if (!gate) return;
|
||||
|
||||
clearCleanupTimer(gate);
|
||||
for (const entry of gate.queue) {
|
||||
clearTimeout(entry.timer);
|
||||
entry.reject(new Error("Semaphore reset"));
|
||||
}
|
||||
gates.delete(semaphoreKey);
|
||||
}
|
||||
|
||||
/**
|
||||
* Reset all keys and reject queued waiters.
|
||||
*/
|
||||
export function resetAll(): void {
|
||||
for (const key of gates.keys()) {
|
||||
reset(key);
|
||||
}
|
||||
}
|
||||
@@ -140,8 +140,10 @@ const STRATEGY_GUIDANCE_FALLBACK = {
|
||||
const ADVANCED_FIELD_HELP_FALLBACK = {
|
||||
maxRetries: "How many retries are attempted before failing the request.",
|
||||
retryDelay: "Initial delay between retries. Higher values reduce burst pressure.",
|
||||
concurrencyPerModel: "Max simultaneous requests sent to each model in round-robin.",
|
||||
queueTimeout: "How long a request can wait in queue before timeout in round-robin.",
|
||||
concurrencyPerModel:
|
||||
"Round-robin combo/model limit: max simultaneous requests sent to each model target. This is separate from any provider account-only cap.",
|
||||
queueTimeout:
|
||||
"How long a request can wait for a round-robin model slot before timing out. This queue is separate from any account-only concurrency cap.",
|
||||
};
|
||||
|
||||
const LEGACY_COMBO_RESILIENCE_KEYS = new Set([
|
||||
|
||||
@@ -543,6 +543,7 @@ interface EditConnectionModalConnection {
|
||||
name?: string;
|
||||
email?: string;
|
||||
priority?: number;
|
||||
maxConcurrent?: number | null;
|
||||
authType?: string;
|
||||
provider?: string;
|
||||
providerSpecificData?: Record<string, unknown>;
|
||||
@@ -5910,6 +5911,7 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec
|
||||
const [formData, setFormData] = useState({
|
||||
name: "",
|
||||
priority: 1,
|
||||
maxConcurrent: "",
|
||||
apiKey: "",
|
||||
healthCheckInterval: 60,
|
||||
baseUrl: "",
|
||||
@@ -5980,6 +5982,10 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec
|
||||
setFormData({
|
||||
name: connection.name || "",
|
||||
priority: connection.priority || 1,
|
||||
maxConcurrent:
|
||||
connection.maxConcurrent === null || connection.maxConcurrent === undefined
|
||||
? ""
|
||||
: String(connection.maxConcurrent),
|
||||
apiKey: "",
|
||||
healthCheckInterval: connection.healthCheckInterval ?? 60,
|
||||
baseUrl: existingBaseUrl || defaultBaseUrl,
|
||||
@@ -6077,9 +6083,21 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec
|
||||
setSaving(true);
|
||||
setSaveError(null);
|
||||
try {
|
||||
const trimmedMaxConcurrent = formData.maxConcurrent.trim();
|
||||
let parsedMaxConcurrent: number | null = null;
|
||||
if (trimmedMaxConcurrent) {
|
||||
const numericMaxConcurrent = Number(trimmedMaxConcurrent);
|
||||
if (!Number.isInteger(numericMaxConcurrent) || numericMaxConcurrent < 0) {
|
||||
setSaveError("Max concurrent must be a whole number greater than or equal to 0.");
|
||||
return;
|
||||
}
|
||||
parsedMaxConcurrent = numericMaxConcurrent;
|
||||
}
|
||||
|
||||
const updates: any = {
|
||||
name: formData.name,
|
||||
priority: formData.priority,
|
||||
maxConcurrent: parsedMaxConcurrent,
|
||||
healthCheckInterval: formData.healthCheckInterval,
|
||||
};
|
||||
|
||||
@@ -6340,6 +6358,30 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec
|
||||
setFormData({ ...formData, priority: Number.parseInt(e.target.value) || 1 })
|
||||
}
|
||||
/>
|
||||
<Input
|
||||
label={t("accountConcurrencyCapLabel")}
|
||||
type="number"
|
||||
min={0}
|
||||
step={1}
|
||||
value={formData.maxConcurrent}
|
||||
onChange={(e) => {
|
||||
const nextValue = e.target.value;
|
||||
setFormData({ ...formData, maxConcurrent: nextValue });
|
||||
if (saveError && nextValue.trim()) {
|
||||
const numericValue = Number(nextValue);
|
||||
if (Number.isInteger(numericValue) && numericValue >= 0) {
|
||||
setSaveError(null);
|
||||
}
|
||||
}
|
||||
}}
|
||||
placeholder="0"
|
||||
hint={t("accountConcurrencyCapHint")}
|
||||
/>
|
||||
{saveError && (
|
||||
<div className="text-sm text-red-500 bg-red-500/10 border border-red-500/20 rounded-lg px-3 py-2">
|
||||
{saveError}
|
||||
</div>
|
||||
)}
|
||||
{!isOAuth && (
|
||||
<>
|
||||
<div className="flex gap-2">
|
||||
@@ -6385,11 +6427,6 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec
|
||||
{validationResult === "success" ? t("valid") : t("invalid")}
|
||||
</Badge>
|
||||
)}
|
||||
{saveError && (
|
||||
<div className="text-sm text-red-500 bg-red-500/10 border border-red-500/20 rounded-lg px-3 py-2">
|
||||
{saveError}
|
||||
</div>
|
||||
)}
|
||||
<button
|
||||
type="button"
|
||||
className="text-sm text-text-muted hover:text-text-primary flex items-center gap-1"
|
||||
|
||||
@@ -1907,6 +1907,8 @@
|
||||
"save": "Save",
|
||||
"editConnection": "Edit Connection",
|
||||
"accountName": "Account name",
|
||||
"accountConcurrencyCapLabel": "Provider account max concurrent requests",
|
||||
"accountConcurrencyCapHint": "Maximum simultaneous requests for this provider account. Leave empty or set 0 for no limit. Helps prevent limits enforced by the provider itself before requests fail.",
|
||||
"email": "Email",
|
||||
"healthCheckMinutes": "Health Check (min)",
|
||||
"healthCheckHint": "Proactive token refresh interval. 0 = disabled.",
|
||||
@@ -2314,6 +2316,8 @@
|
||||
"defaultSafetyNet": "Default Safety Net",
|
||||
"rpm": "RPM",
|
||||
"minGap": "Min Gap",
|
||||
"comboConcurrencyLabel": "Combo round-robin max concurrent requests per model",
|
||||
"comboConcurrencyHint": "Maximum simultaneous requests per model during combo round-robin. This setting only applies to combo strategies.",
|
||||
"maxConcurrent": "Max Concurrent",
|
||||
"activeLimiters": "Active Limiters",
|
||||
"noActiveLimiters": "No active rate limiters yet.",
|
||||
|
||||
@@ -1846,6 +1846,8 @@
|
||||
"save": "저장",
|
||||
"editConnection": "연결 편집",
|
||||
"accountName": "계정 이름",
|
||||
"accountConcurrencyCapLabel": "Provider 계정 최대 동시 요청 수",
|
||||
"accountConcurrencyCapHint": "이 provider 계정의 최대 동시 요청 수입니다. 비워 두거나 0으로 설정하면 제한이 없습니다. provider 자체가 부과하는 제한에 먼저 걸리지 않도록 돕습니다.",
|
||||
"email": "이메일",
|
||||
"healthCheckMinutes": "상태 점검(분)",
|
||||
"healthCheckHint": "사전 토큰 새로 고침 간격. 0 = 비활성화됨.",
|
||||
@@ -2216,6 +2218,8 @@
|
||||
"defaultSafetyNet": "기본 안전망",
|
||||
"rpm": "RPM",
|
||||
"minGap": "민갭",
|
||||
"comboConcurrencyLabel": "콤보 라운드로빈 모델별 최대 동시 요청 수",
|
||||
"comboConcurrencyHint": "Combo 라운드로빈에서 모델별 최대 동시 요청 수입니다. 이 설정은 combo 전략에만 적용됩니다.",
|
||||
"maxConcurrent": "최대 동시",
|
||||
"activeLimiters": "활성 리미터",
|
||||
"noActiveLimiters": "아직 활성 속도 제한기가 없습니다.",
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
-- 028_provider_connection_max_concurrent.sql
|
||||
-- Adds account-native concurrency cap to provider_connections.
|
||||
-- This defines the maximum concurrent requests allowed for a specific account (connection).
|
||||
-- Coexists with existing combo-level concurrencyPerModel, which is separate.
|
||||
|
||||
-- Add max_concurrent column (NULL = unlimited, uses provider defaults or combo rules)
|
||||
ALTER TABLE provider_connections ADD COLUMN max_concurrent INTEGER;
|
||||
|
||||
-- Index for provider-level filtering
|
||||
CREATE INDEX IF NOT EXISTS idx_pc_max_concurrent ON provider_connections(provider, max_concurrent);
|
||||
@@ -21,6 +21,26 @@ interface DbLike {
|
||||
prepare: <TRow = unknown>(sql: string) => StatementLike<TRow>;
|
||||
}
|
||||
|
||||
function withNullableMaxConcurrent(
|
||||
record: JsonRecord,
|
||||
source: JsonRecord | null | undefined
|
||||
): JsonRecord {
|
||||
if (!source || !Object.hasOwn(source, "maxConcurrent")) {
|
||||
return record;
|
||||
}
|
||||
|
||||
const sourceMaxConcurrent = source.maxConcurrent;
|
||||
const normalizedMaxConcurrent =
|
||||
typeof sourceMaxConcurrent === "number" || sourceMaxConcurrent === null
|
||||
? sourceMaxConcurrent
|
||||
: record.maxConcurrent;
|
||||
|
||||
return {
|
||||
...record,
|
||||
maxConcurrent: normalizedMaxConcurrent,
|
||||
};
|
||||
}
|
||||
|
||||
function toRecord(value: unknown): JsonRecord {
|
||||
return value && typeof value === "object" ? (value as JsonRecord) : {};
|
||||
}
|
||||
@@ -56,13 +76,19 @@ export async function getProviderConnections(filter: JsonRecord = {}) {
|
||||
sql += " ORDER BY priority ASC, updated_at DESC";
|
||||
|
||||
const rows = db.prepare(sql).all(params);
|
||||
return rows.map((r) => decryptConnectionFields(cleanNulls(rowToCamel(r))));
|
||||
return rows.map((r) => {
|
||||
const camelRow = rowToCamel(r);
|
||||
return decryptConnectionFields(withNullableMaxConcurrent(cleanNulls(camelRow), camelRow));
|
||||
});
|
||||
}
|
||||
|
||||
export async function getProviderConnectionById(id: string) {
|
||||
const db = getDbInstance() as unknown as DbLike;
|
||||
const row = db.prepare("SELECT * FROM provider_connections WHERE id = ?").get(id);
|
||||
return row ? decryptConnectionFields(cleanNulls(rowToCamel(row))) : null;
|
||||
if (!row) return null;
|
||||
|
||||
const camelRow = rowToCamel(row);
|
||||
return decryptConnectionFields(withNullableMaxConcurrent(cleanNulls(camelRow), camelRow));
|
||||
}
|
||||
|
||||
export async function createProviderConnection(data: JsonRecord) {
|
||||
@@ -133,7 +159,7 @@ export async function createProviderConnection(data: JsonRecord) {
|
||||
);
|
||||
_updateConnectionRow(db, existingId, merged);
|
||||
backupDbFile("pre-write");
|
||||
return cleanNulls(merged);
|
||||
return withNullableMaxConcurrent(cleanNulls(merged), merged);
|
||||
}
|
||||
|
||||
// Generate name: prefer explicit name, then email, then a stable short-ID label.
|
||||
@@ -195,6 +221,7 @@ export async function createProviderConnection(data: JsonRecord) {
|
||||
"consecutiveUseCount",
|
||||
"rateLimitProtection",
|
||||
"group",
|
||||
"maxConcurrent",
|
||||
];
|
||||
for (const field of optionalFields) {
|
||||
if (data[field] !== undefined && data[field] !== null) {
|
||||
@@ -213,7 +240,7 @@ export async function createProviderConnection(data: JsonRecord) {
|
||||
backupDbFile("pre-write");
|
||||
invalidateDbCache("connections"); // Bust connections read cache
|
||||
|
||||
return cleanNulls(connection);
|
||||
return withNullableMaxConcurrent(cleanNulls(connection), connection);
|
||||
}
|
||||
|
||||
function _insertConnectionRow(db: DbLike, conn: JsonRecord) {
|
||||
@@ -227,7 +254,8 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) {
|
||||
rate_limited_until, health_check_interval, last_health_check_at,
|
||||
last_tested, api_key, id_token, provider_specific_data,
|
||||
expires_in, display_name, global_priority, default_model,
|
||||
token_type, consecutive_use_count, rate_limit_protection, last_used_at, "group", created_at, updated_at
|
||||
token_type, consecutive_use_count, rate_limit_protection, last_used_at, "group", max_concurrent,
|
||||
created_at, updated_at
|
||||
) VALUES (
|
||||
@id, @provider, @authType, @name, @email, @priority, @isActive,
|
||||
@accessToken, @refreshToken, @expiresAt, @tokenExpiresAt,
|
||||
@@ -236,7 +264,8 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) {
|
||||
@rateLimitedUntil, @healthCheckInterval, @lastHealthCheckAt,
|
||||
@lastTested, @apiKey, @idToken, @providerSpecificData,
|
||||
@expiresIn, @displayName, @globalPriority, @defaultModel,
|
||||
@tokenType, @consecutiveUseCount, @rateLimitProtection, @lastUsedAt, @group, @createdAt, @updatedAt
|
||||
@tokenType, @consecutiveUseCount, @rateLimitProtection, @lastUsedAt, @group, @maxConcurrent,
|
||||
@createdAt, @updatedAt
|
||||
)
|
||||
`
|
||||
).run({
|
||||
@@ -279,6 +308,7 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) {
|
||||
conn.rateLimitProtection === true || conn.rateLimitProtection === 1 ? 1 : 0,
|
||||
lastUsedAt: conn.lastUsedAt || null,
|
||||
group: conn.group || null,
|
||||
maxConcurrent: conn.maxConcurrent ?? null,
|
||||
createdAt: conn.createdAt,
|
||||
updatedAt: conn.updatedAt,
|
||||
});
|
||||
@@ -304,6 +334,7 @@ function _updateConnectionRow(db: DbLike, id: string, data: JsonRecord) {
|
||||
rate_limit_protection = @rateLimitProtection,
|
||||
last_used_at = @lastUsedAt,
|
||||
"group" = @group,
|
||||
max_concurrent = @maxConcurrent,
|
||||
updated_at = @updatedAt
|
||||
WHERE id = @id
|
||||
`
|
||||
@@ -347,6 +378,7 @@ function _updateConnectionRow(db: DbLike, id: string, data: JsonRecord) {
|
||||
data.rateLimitProtection === true || data.rateLimitProtection === 1 ? 1 : 0,
|
||||
lastUsedAt: data.lastUsedAt || null,
|
||||
group: data.group || null,
|
||||
maxConcurrent: data.maxConcurrent ?? null,
|
||||
updatedAt: now,
|
||||
});
|
||||
}
|
||||
@@ -378,7 +410,7 @@ export async function updateProviderConnection(id: string, data: JsonRecord) {
|
||||
_reorderConnections(db, providerId);
|
||||
}
|
||||
|
||||
return cleanNulls(merged);
|
||||
return withNullableMaxConcurrent(cleanNulls(merged), merged);
|
||||
}
|
||||
|
||||
export async function deleteProviderConnection(id: string) {
|
||||
|
||||
@@ -1477,6 +1477,7 @@ export const updateProviderConnectionSchema = z
|
||||
lastTested: z.union([z.string(), z.null()]).optional(),
|
||||
healthCheckInterval: z.coerce.number().int().min(0).optional(),
|
||||
group: z.union([z.string().max(100), z.null()]).optional(),
|
||||
maxConcurrent: z.union([z.null(), z.coerce.number().int().min(0)]).optional(),
|
||||
// Partial patch of per-connection provider-specific settings (e.g. quota toggles)
|
||||
providerSpecificData: z
|
||||
.record(z.string(), z.unknown())
|
||||
|
||||
116
tests/unit/account-concurrency-cap.test.ts
Normal file
116
tests/unit/account-concurrency-cap.test.ts
Normal file
@@ -0,0 +1,116 @@
|
||||
import { after, beforeEach, describe, it } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
|
||||
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-account-cap-"));
|
||||
process.env.DATA_DIR = TEST_DATA_DIR;
|
||||
|
||||
const core = await import("../../src/lib/db/core.ts");
|
||||
const providersDb = await import("../../src/lib/db/providers.ts");
|
||||
const { updateProviderConnectionSchema } = await import("../../src/shared/validation/schemas.ts");
|
||||
|
||||
type Connection = Awaited<ReturnType<typeof providersDb.createProviderConnection>>;
|
||||
type StoredConnection = Awaited<ReturnType<typeof providersDb.getProviderConnectionById>>;
|
||||
|
||||
function assertConnection(connection: Connection): asserts connection is NonNullable<Connection> {
|
||||
assert.ok(connection);
|
||||
}
|
||||
|
||||
function assertStoredConnection(
|
||||
connection: StoredConnection
|
||||
): asserts connection is NonNullable<StoredConnection> {
|
||||
assert.ok(connection);
|
||||
}
|
||||
|
||||
function getConnectionId(connection: NonNullable<Connection>): string {
|
||||
assert.equal(typeof connection.id, "string");
|
||||
return connection.id as string;
|
||||
}
|
||||
|
||||
async function resetStorage() {
|
||||
core.resetDbInstance();
|
||||
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
|
||||
fs.mkdirSync(TEST_DATA_DIR, { recursive: true });
|
||||
}
|
||||
|
||||
async function createConnection(maxConcurrent: number | null): Promise<Connection> {
|
||||
return providersDb.createProviderConnection({
|
||||
provider: "openai",
|
||||
authType: "apikey",
|
||||
name: `openai-${String(maxConcurrent)}-${Math.random().toString(16).slice(2, 8)}`,
|
||||
apiKey: "sk-test",
|
||||
maxConcurrent,
|
||||
});
|
||||
}
|
||||
|
||||
beforeEach(async () => {
|
||||
await resetStorage();
|
||||
});
|
||||
|
||||
after(() => {
|
||||
core.resetDbInstance();
|
||||
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
describe("maxConcurrent DB round-trip", () => {
|
||||
it("stores and retrieves maxConcurrent=3", async () => {
|
||||
const created = await createConnection(3);
|
||||
assertConnection(created);
|
||||
const connectionId = getConnectionId(created);
|
||||
|
||||
const stored = await providersDb.getProviderConnectionById(connectionId);
|
||||
assertStoredConnection(stored);
|
||||
assert.equal(stored.maxConcurrent, 3);
|
||||
});
|
||||
|
||||
it("stores and retrieves maxConcurrent=null", async () => {
|
||||
const created = await createConnection(null);
|
||||
assertConnection(created);
|
||||
const connectionId = getConnectionId(created);
|
||||
|
||||
const stored = await providersDb.getProviderConnectionById(connectionId);
|
||||
assertStoredConnection(stored);
|
||||
assert.equal(stored.maxConcurrent, null);
|
||||
});
|
||||
|
||||
it("stores and retrieves maxConcurrent=0", async () => {
|
||||
const created = await createConnection(3);
|
||||
assertConnection(created);
|
||||
const connectionId = getConnectionId(created);
|
||||
|
||||
const updated = await providersDb.updateProviderConnection(connectionId, { maxConcurrent: 0 });
|
||||
assertConnection(updated);
|
||||
assert.equal(updated.maxConcurrent, 0);
|
||||
|
||||
const stored = await providersDb.getProviderConnectionById(connectionId);
|
||||
assertStoredConnection(stored);
|
||||
assert.equal(stored.maxConcurrent, 0);
|
||||
});
|
||||
});
|
||||
|
||||
describe("maxConcurrent validation", () => {
|
||||
it("rejects negative values", () => {
|
||||
const result = updateProviderConnectionSchema.safeParse({ maxConcurrent: -1 });
|
||||
|
||||
assert.equal(result.success, false);
|
||||
});
|
||||
|
||||
it("accepts positive integers", () => {
|
||||
const result = updateProviderConnectionSchema.safeParse({ maxConcurrent: 3 });
|
||||
|
||||
assert.equal(result.success, true);
|
||||
assert.equal(result.data.maxConcurrent, 3);
|
||||
});
|
||||
|
||||
it("accepts null/undefined as unlimited", () => {
|
||||
const nullResult = updateProviderConnectionSchema.safeParse({ maxConcurrent: null });
|
||||
const undefinedResult = updateProviderConnectionSchema.safeParse({ name: "unchanged" });
|
||||
|
||||
assert.equal(nullResult.success, true);
|
||||
assert.equal(nullResult.data.maxConcurrent, null);
|
||||
assert.equal(undefinedResult.success, true);
|
||||
assert.equal(undefinedResult.data.maxConcurrent, undefined);
|
||||
});
|
||||
});
|
||||
175
tests/unit/accountSemaphore.test.ts
Normal file
175
tests/unit/accountSemaphore.test.ts
Normal file
@@ -0,0 +1,175 @@
|
||||
import assert from "node:assert/strict";
|
||||
import { afterEach, describe, it } from "node:test";
|
||||
|
||||
import {
|
||||
acquire,
|
||||
buildAccountSemaphoreKey,
|
||||
getStats,
|
||||
markBlocked,
|
||||
reset,
|
||||
resetAll,
|
||||
} from "../../open-sse/services/accountSemaphore";
|
||||
|
||||
afterEach(() => {
|
||||
resetAll();
|
||||
});
|
||||
|
||||
describe("accountSemaphore", async () => {
|
||||
it("queues requests beyond the account cap and drains on release", async () => {
|
||||
const key = buildAccountSemaphoreKey({
|
||||
provider: "alibaba",
|
||||
accountKey: "acct-1",
|
||||
});
|
||||
|
||||
const releaseA = await acquire(key, { maxConcurrency: 2, timeoutMs: 200 });
|
||||
const releaseB = await acquire(key, { maxConcurrency: 2, timeoutMs: 200 });
|
||||
const queued = acquire(key, { maxConcurrency: 2, timeoutMs: 200 });
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
|
||||
assert.deepEqual(getStats()[key], {
|
||||
running: 2,
|
||||
queued: 1,
|
||||
maxConcurrency: 2,
|
||||
blockedUntil: null,
|
||||
});
|
||||
|
||||
releaseA();
|
||||
const releaseC = await queued;
|
||||
|
||||
assert.deepEqual(getStats()[key], {
|
||||
running: 2,
|
||||
queued: 0,
|
||||
maxConcurrency: 2,
|
||||
blockedUntil: null,
|
||||
});
|
||||
|
||||
releaseA();
|
||||
releaseB();
|
||||
releaseC();
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
assert.equal(getStats()[key], undefined);
|
||||
});
|
||||
|
||||
it("returns a no-op release when concurrency is bypassed", async () => {
|
||||
const key = buildAccountSemaphoreKey({
|
||||
provider: "alibaba",
|
||||
accountKey: "acct-bypass",
|
||||
});
|
||||
|
||||
const release = await acquire(key, { maxConcurrency: 0, timeoutMs: 50 });
|
||||
|
||||
assert.deepEqual(getStats(), {});
|
||||
|
||||
release();
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
assert.equal(getStats()[key], undefined);
|
||||
});
|
||||
|
||||
it("uses SEMAPHORE_TIMEOUT for timed out queued requests", async () => {
|
||||
const key = buildAccountSemaphoreKey({
|
||||
provider: "alibaba",
|
||||
accountKey: "acct-timeout",
|
||||
});
|
||||
|
||||
const releaseA = await acquire(key, { maxConcurrency: 1, timeoutMs: 200 });
|
||||
const queued = acquire(key, { maxConcurrency: 1, timeoutMs: 200 });
|
||||
|
||||
try {
|
||||
await queued;
|
||||
assert.fail("Expected timeout error");
|
||||
} catch (err: unknown) {
|
||||
assert.ok(err instanceof Error);
|
||||
const error = err as Error & { code?: string };
|
||||
assert.equal(error.code, "SEMAPHORE_TIMEOUT");
|
||||
}
|
||||
|
||||
releaseA();
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
assert.equal(getStats()[key], undefined);
|
||||
});
|
||||
|
||||
it("keeps release idempotent for finally blocks", async () => {
|
||||
const key = buildAccountSemaphoreKey({
|
||||
provider: "alibaba",
|
||||
accountKey: "acct-idempotent",
|
||||
});
|
||||
|
||||
const releaseA = await acquire(key, { maxConcurrency: 1, timeoutMs: 200 });
|
||||
|
||||
// Simulate a finally block calling release twice
|
||||
releaseA();
|
||||
releaseA();
|
||||
releaseA();
|
||||
|
||||
// The second acquire should succeed immediately (slot was released)
|
||||
const releaseB = await acquire(key, { maxConcurrency: 1, timeoutMs: 200 });
|
||||
|
||||
assert.deepEqual(getStats()[key], {
|
||||
running: 1,
|
||||
queued: 0,
|
||||
maxConcurrency: 1,
|
||||
blockedUntil: null,
|
||||
});
|
||||
|
||||
releaseB();
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
assert.equal(getStats()[key], undefined);
|
||||
});
|
||||
|
||||
it("supports temporary blocking and explicit reset hooks", async () => {
|
||||
const key = buildAccountSemaphoreKey({
|
||||
provider: "alibaba",
|
||||
accountKey: "acct-blocked",
|
||||
});
|
||||
|
||||
await acquire(key, { maxConcurrency: 1, timeoutMs: 200 });
|
||||
|
||||
assert.deepEqual(getStats()[key], {
|
||||
running: 1,
|
||||
queued: 0,
|
||||
maxConcurrency: 1,
|
||||
blockedUntil: null,
|
||||
});
|
||||
|
||||
markBlocked(key, Date.now() + 50);
|
||||
|
||||
// Should block even though slot is available
|
||||
const acquired = acquire(key, { maxConcurrency: 1, timeoutMs: 100 });
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
|
||||
// Should still be queued because the gate is blocked
|
||||
const stats = getStats()[key];
|
||||
assert.equal(stats.running, 1);
|
||||
assert.equal(stats.queued, 1);
|
||||
assert.equal(stats.maxConcurrency, 1);
|
||||
assert.ok(stats.blockedUntil !== null, "blockedUntil should be set");
|
||||
|
||||
reset(key);
|
||||
|
||||
await assert.rejects(async () => {
|
||||
await acquired;
|
||||
});
|
||||
});
|
||||
|
||||
it("preserves existing maxConcurrency when markBlocked is applied", async () => {
|
||||
const key = buildAccountSemaphoreKey({
|
||||
provider: "alibaba",
|
||||
accountKey: "acct-preserve",
|
||||
});
|
||||
|
||||
await acquire(key, { maxConcurrency: 2, timeoutMs: 200 });
|
||||
markBlocked(key, Date.now() + 50);
|
||||
|
||||
const stats = getStats()[key];
|
||||
assert.equal(stats.running, 1);
|
||||
assert.equal(stats.queued, 0);
|
||||
assert.equal(stats.maxConcurrency, 2);
|
||||
assert.ok(stats.blockedUntil !== null, "blockedUntil should be set");
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user