diff --git a/packages/lib/RateLimiter.js b/packages/lib/RateLimiter.js new file mode 100644 index 00000000..69adf335 --- /dev/null +++ b/packages/lib/RateLimiter.js @@ -0,0 +1,104 @@ +"use strict"; + + +/** + * Helper class for rate limiting actions + * + * Provides a simple interface for implementing actions that are rate + * limited on invocations. + * + * @example + * const rateLimiter = RateLimiter({ + * action: () => { console.log("action"); }, + * maxRate: 1, + * }); + * + * for (let i = 0; i < 100; i++) { + * rateLimiter.activate(); + * // Prints action twice to the console, once for the first activation + * // and one more time 1 second later for 99 other activations that were + * // exceeding the rate limit. + * } + * @alias module:lib/RateLimiter + */ +class RateLimiter { + /** + * Construct RateLimiter helper + * + * The options passed may also be modified on the instance at runtime. + * + * @param {object=} options - Options for this rate limiter + * @param {number=} options.maxRate - + * Maximum rate of activations per second. + * @param {function()=} options.action - + * Callback invoked at maste `maxRate` per second on activation. + */ + constructor(options = {}) { + /* + * Maximum rate of activations per second. + * @type {number} + */ + this.maxRate = options.maxRate || 1; + + /** + * Callback invoked at most `maxRate` per second on activation. + * @type {function()} + */ + this.action = options.action || null; + + this._lastRun = 0; + this._runTimeout = null; + this.lastProgressBroadcast = Date.now(); + } + + /** + * Activate the rate limited action + * + * Checks if the invocation of the rate limited action now would fall + * within the rate limit. If it is within the limit then the action is + * invoked immediatly, otherwise a timeout is registered to invoke the + * action at next point in time where it will be within the rate limit. + * + * @returns {boolean} true if activation was within the rate limit. + */ + activate() { + let now = Date.now(); + if (now < this._lastRun + 1000 / this.maxRate) { + if (!this._runTimeout && this.action !== null) { + this._runTimeout = setTimeout(() => { + this._runTimeout = null; + this._lastRun = Date.now(); + this.action(); + }, this._lastRun + 1000 / this.maxRate - now); + } + return false; + } + + // We may end up activating before the timeout does. + if (this._runTimeout) { + clearTimeout(this._runTimeout); + this._runTimeout = null; + } + + this._lastRun = now; + if (this.action !== null) { + this.action(); + } + return true; + } + + /** + * Cancel pending action + * + * Cancel pending action invocation. This is necessary when shutting + * down the use of the rate limit as a timeout is registered to invoke + * the next action when activated faster than the rate limit allows. + */ + cancel() { + if (this._runTimeout) { + clearTimeout(this._runTimeout); + } + } +} + +module.exports = RateLimiter; diff --git a/plugins/research_sync/master.js b/plugins/research_sync/master.js index 1475680a..9d735ceb 100644 --- a/plugins/research_sync/master.js +++ b/plugins/research_sync/master.js @@ -3,41 +3,27 @@ const fs = require("fs-extra"); const path = require("path"); const libPlugin = require("@clusterio/lib/plugin"); +const RateLimiter = require("@clusterio/lib/RateLimiter"); class MasterPlugin extends libPlugin.BaseMasterPlugin { async init() { this.technologies = await loadTechnologies(this.master.config, this.logger); + this.progressRateLimiter = new RateLimiter({ + maxRate: 1, + action: () => this.broadcastProgress(), + }); + this.lastProgressBroadcast = Date.now(); this.progressBroadcastId = null; this.progressToBroadcast = new Set(); } async onShutdown() { - if (this.progressBroadcastId) { - clearTimeout(this.progressBroadcastId); - } + this.progressRateLimiter.cancel(); await saveTechnologies(this.master.config, this.technologies, this.logger); } - registerProgress() { - // Rate limit progress broadcasts to one per second - if (Date.now() < this.lastProgressBroadcast + 1000) { - if (!this.progressBroadcastId) { - this.progressBroadcastId = setTimeout(() => { - this.broadcastProgress(); - this.progressBroadcastId = null; - }, 2000); - } - - } else { - if (this.progressBroadcastId) { - clearTimeout(this.progressBroadcastId); - } - this.broadcastProgress(); - } - } - broadcastProgress() { let techs = []; for (let name of this.progressToBroadcast) { @@ -81,7 +67,7 @@ class MasterPlugin extends libPlugin.BaseMasterPlugin { if (newProgress < 1) { tech.progress = newProgress; this.progressToBroadcast.add(name); - this.registerProgress(); + this.progressRateLimiter.activate(); } else { tech.researched = true; @@ -148,7 +134,7 @@ class MasterPlugin extends libPlugin.BaseMasterPlugin { } } } - this.registerProgress(); + this.progressRateLimiter.activate(); let technologies = []; for (let [name, tech] of this.technologies) { diff --git a/test/lib/RateLimiter.js b/test/lib/RateLimiter.js new file mode 100644 index 00000000..d7448d3f --- /dev/null +++ b/test/lib/RateLimiter.js @@ -0,0 +1,64 @@ +"use strict"; +const assert = require("assert").strict; + +const { wait } = require("@clusterio/lib/helpers"); +const RateLimiter = require("@clusterio/lib/RateLimiter"); + + +describe("lib/RateLimiter", function() { + describe("class RateLimiter", function() { + let rateLimiter; + beforeEach(function() { + rateLimiter = new RateLimiter(); + }); + afterEach(function() { + rateLimiter.cancel(); + }); + + it("should rate limit activations", function() { + assert(rateLimiter.activate(), "fist invocation was rate limited"); + assert(!rateLimiter.activate(), "second invocation was not rate limited"); + }); + + it("should invoke action immediately when not rate limited", function() { + let invoked = false; + rateLimiter.action = () => { invoked = true; }; + rateLimiter.activate(); + assert(invoked, "action was not invoked"); + }); + + it("should invoke action at a later point in time if rate limited", async function() { + let invokedCount = 0; + rateLimiter.maxRate = 100; + rateLimiter.action = () => { invokedCount += 1; }; + rateLimiter.activate(); + rateLimiter.activate(); + assert.equal(invokedCount, 1, "expected one invocation"); + await wait(20); + assert.equal(invokedCount, 2, "expected two invocations"); + }); + + it("should invoke action before timeout if maxRate changed", async function() { + let invokedCount = 0; + rateLimiter.action = () => { invokedCount += 1; }; + rateLimiter.activate(); + rateLimiter.activate(); + assert.equal(invokedCount, 1, "expected one invocation"); + await wait(20); + rateLimiter.maxRate = 100; + rateLimiter.activate(); + assert.equal(invokedCount, 2, "expected two invocations"); + }); + + it("should cancel action if canceled before timeout", async function() { + let invokedCount = 0; + rateLimiter.action = () => { invokedCount += 1; }; + rateLimiter.activate(); + rateLimiter.activate(); + rateLimiter.cancel(); + assert.equal(invokedCount, 1, "expected one invocation"); + await wait(20); + assert.equal(invokedCount, 1, "expected one invocation"); + }); + }); +});