81 lines
2.3 KiB
JavaScript
81 lines
2.3 KiB
JavaScript
class ToolRateLimits {
|
|
constructor(options = {}) {
|
|
this.now = options.now || Date.now;
|
|
this.buckets = new Map();
|
|
}
|
|
|
|
consume({ actor, origin, server }, settings) {
|
|
const checks = [
|
|
[`user:${actor}`, settings.per_user_per_minute],
|
|
[`origin:${origin}`, settings.per_origin_per_minute],
|
|
[`server:${origin}:${server}`, settings.per_server_per_minute]
|
|
];
|
|
let retryAfter = 0;
|
|
for (const [key, maximum] of checks) {
|
|
const result = this.inspect(key, maximum);
|
|
if (!result.allowed) retryAfter = Math.max(retryAfter, result.retry_after_seconds);
|
|
}
|
|
if (retryAfter) return { allowed: false, retry_after_seconds: retryAfter };
|
|
for (const [key] of checks) this.record(key);
|
|
return { allowed: true, retry_after_seconds: 0 };
|
|
}
|
|
|
|
inspect(key, maximum) {
|
|
const cutoff = this.now() - 60000;
|
|
const recent = (this.buckets.get(key) || []).filter((timestamp) => timestamp > cutoff);
|
|
this.buckets.set(key, recent);
|
|
if (recent.length < maximum) return { allowed: true, retry_after_seconds: 0 };
|
|
return {
|
|
allowed: false,
|
|
retry_after_seconds: Math.max(1, Math.ceil((recent[0] + 60000 - this.now()) / 1000))
|
|
};
|
|
}
|
|
|
|
record(key) {
|
|
const recent = this.buckets.get(key) || [];
|
|
recent.push(this.now());
|
|
this.buckets.set(key, recent);
|
|
}
|
|
}
|
|
|
|
class ToolConcurrency {
|
|
constructor(maximum = 3, maximumQueue = 20) {
|
|
this.maximum = Math.max(1, Number(maximum) || 3);
|
|
this.maximumQueue = Math.max(0, Number(maximumQueue) || 20);
|
|
this.active = 0;
|
|
this.queue = [];
|
|
}
|
|
|
|
run(callback) {
|
|
if (this.active < this.maximum) return this.start(callback);
|
|
if (this.queue.length >= this.maximumQueue) {
|
|
return Promise.reject(Object.assign(
|
|
new Error("Web lookup concurrency queue is full."),
|
|
{ code: "concurrency_limited" }
|
|
));
|
|
}
|
|
return new Promise((resolve, reject) => {
|
|
this.queue.push({ callback, resolve, reject });
|
|
});
|
|
}
|
|
|
|
async start(callback) {
|
|
this.active += 1;
|
|
try {
|
|
return await callback();
|
|
} finally {
|
|
this.active -= 1;
|
|
this.drain();
|
|
}
|
|
}
|
|
|
|
drain() {
|
|
while (this.active < this.maximum && this.queue.length) {
|
|
const next = this.queue.shift();
|
|
this.start(next.callback).then(next.resolve, next.reject);
|
|
}
|
|
}
|
|
}
|
|
|
|
module.exports = { ToolConcurrency, ToolRateLimits };
|