97 lines
3.2 KiB
JavaScript
97 lines
3.2 KiB
JavaScript
const assert = require("node:assert/strict");
|
|
const { createRequire } = require("node:module");
|
|
|
|
const nodeRequire = createRequire(__filename);
|
|
|
|
function loadLimiterWithPool(pool) {
|
|
const dbPath = nodeRequire.resolve("../src/db");
|
|
const limiterPath = nodeRequire.resolve("../src/providerPollLimiter");
|
|
const originalDbModule = nodeRequire.cache[dbPath];
|
|
const originalLimiterModule = nodeRequire.cache[limiterPath];
|
|
|
|
delete nodeRequire.cache[limiterPath];
|
|
nodeRequire.cache[dbPath] = {
|
|
id: dbPath,
|
|
filename: dbPath,
|
|
loaded: true,
|
|
exports: { pool },
|
|
};
|
|
|
|
return {
|
|
limiter: nodeRequire("../src/providerPollLimiter"),
|
|
restore() {
|
|
delete nodeRequire.cache[limiterPath];
|
|
if (originalLimiterModule) nodeRequire.cache[limiterPath] = originalLimiterModule;
|
|
if (originalDbModule) nodeRequire.cache[dbPath] = originalDbModule;
|
|
else delete nodeRequire.cache[dbPath];
|
|
},
|
|
};
|
|
}
|
|
|
|
function createPool(options = {}) {
|
|
const calls = [];
|
|
return {
|
|
calls,
|
|
async query(sql, params = []) {
|
|
calls.push({ sql, params });
|
|
if (/CREATE TABLE IF NOT EXISTS generation_provider_poll_slots/i.test(sql)) return { rows: [] };
|
|
if (/WITH candidate AS/i.test(sql)) {
|
|
if (options.noAvailableSlot) return { rows: [] };
|
|
return { rows: [{ scope: params[0], slot_no: 2 }] };
|
|
}
|
|
if (/DELETE FROM generation_provider_poll_slots/i.test(sql)) return { rows: [] };
|
|
throw new Error(`Unexpected SQL: ${sql}`);
|
|
},
|
|
};
|
|
}
|
|
|
|
(async () => {
|
|
const previousLimit = process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY;
|
|
process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY = "3";
|
|
|
|
const pool = createPool();
|
|
const { limiter, restore } = loadLimiterWithPool(pool);
|
|
try {
|
|
const outcome = await limiter.withProviderPollSlot(101, async () => "polled");
|
|
|
|
assert.equal(outcome.acquired, true);
|
|
assert.equal(outcome.value, "polled");
|
|
|
|
const acquireCall = pool.calls.find((call) => /WITH candidate AS/i.test(call.sql));
|
|
assert.equal(acquireCall.params[1], 3);
|
|
assert.equal(acquireCall.params[3], 101);
|
|
|
|
const releaseCall = pool.calls.find((call) => /DELETE FROM generation_provider_poll_slots/i.test(call.sql));
|
|
assert.equal(releaseCall.params[0], acquireCall.params[0]);
|
|
assert.equal(releaseCall.params[1], 2);
|
|
assert.equal(releaseCall.params[2], acquireCall.params[2]);
|
|
} finally {
|
|
if (previousLimit === undefined) delete process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY;
|
|
else process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY = previousLimit;
|
|
restore();
|
|
}
|
|
|
|
const saturatedPool = createPool({ noAvailableSlot: true });
|
|
const { limiter: saturatedLimiter, restore: restoreSaturated } = loadLimiterWithPool(saturatedPool);
|
|
try {
|
|
let called = false;
|
|
const outcome = await saturatedLimiter.withProviderPollSlot(202, async () => {
|
|
called = true;
|
|
return "should-not-run";
|
|
});
|
|
|
|
assert.equal(outcome.acquired, false);
|
|
assert.equal(outcome.value, undefined);
|
|
assert.equal(called, false);
|
|
assert.equal(
|
|
saturatedPool.calls.some((call) => /DELETE FROM generation_provider_poll_slots/i.test(call.sql)),
|
|
false,
|
|
);
|
|
} finally {
|
|
restoreSaturated();
|
|
}
|
|
})().catch((error) => {
|
|
console.error(error);
|
|
process.exitCode = 1;
|
|
});
|