util.throttle: Fix initial time setting (double accounting the first time) and fractional balance updates (0.1*10 was not the same as 1*1)

2 files changed, 133 insertions, 13 deletions

diff --git a/spec/util_throttle_spec.lua b/spec/util_throttle_spec.lua
index a588e262..ed4d810b 100644
--- a/spec/util_throttle_spec.lua
+++ b/spec/util_throttle_spec.lua
@@ -12,19 +12,139 @@ package.loaded["util.time"] = {
 
 local throttle = require "util.throttle";
 
-describe("util.sasl.scram", function()
-	describe("#Hi()", function()
-		it("should work", function()
+describe("util.throttle", function()
+	describe("#create", function()
+		it("should be created with correct values", function()
+			now = 5;
 			local a = throttle.create(3, 10);
+			assert.same(a, { balance = 3, max = 3, rate = 0.3, t = 5 });
 
-			assert.are.equal(a:poll(1), true);  -- 3 -> 2
-			assert.are.equal(a:poll(1), true);  -- 2 -> 1
-			assert.are.equal(a:poll(1), true);  -- 1 -> 0
-			assert.are.equal(a:poll(1), false); -- MEEP, out of credits!
-			later(1);                       -- ... what about
-			assert.are.equal(a:poll(1), false); -- now? - Still no!
-			later(9);                       -- Later that day
-			assert.are.equal(a:poll(1), true);  -- Should be back at 3 credits ... 2
+			local a = throttle.create(3, 5);
+			assert.same(a, { balance = 3, max = 3, rate = 0.6, t = 5 });
+
+			local a = throttle.create(1, 1);
+			assert.same(a, { balance = 1, max = 1, rate = 1, t = 5 });
+
+			local a = throttle.create(10, 10);
+			assert.same(a, { balance = 10, max = 10, rate = 1, t = 5 });
+
+			local a = throttle.create(10, 1);
+			assert.same(a, { balance = 10, max = 10, rate = 10, t = 5 });
+		end);
+	end);
+
+	describe("#update", function()
+		it("does nothing when no time hase passed, even if balance is not full", function()
+			now = 5;
+			local a = throttle.create(10, 10);
+			for i=1,5 do
+				a:update();
+				assert.same(a, { balance = 10, max = 10, rate = 1, t = 5 });
+			end
+			a.balance = 0;
+			for i=1,5 do
+				a:update();
+				assert.same(a, { balance = 0, max = 10, rate = 1, t = 5 });
+			end
+		end);
+		it("updates only time when time passes but balance is full", function()
+			now = 5;
+			local a = throttle.create(10, 10);
+			for i=1,5 do
+				later(5);
+				a:update();
+				assert.same(a, { balance = 10, max = 10, rate = 1, t = 5 + i*5 });
+			end
+		end);
+		it("updates balance when balance has room to grow as time passes", function()
+			now = 5;
+			local a = throttle.create(10, 10);
+			a.balance = 0;
+			assert.same(a, { balance = 0, max = 10, rate = 1, t = 5 });
+
+			later(1);
+			a:update();
+			assert.same(a, { balance = 1, max = 10, rate = 1, t = 6 });
+
+			later(3);
+			a:update();
+			assert.same(a, { balance = 4, max = 10, rate = 1, t = 9 });
+
+			later(10);
+			a:update();
+			assert.same(a, { balance = 10, max = 10, rate = 1, t = 19 });
+		end);
+		it("handles 10 x 0.1s updates the same as 1 x 1s update ", function()
+			now = 5;
+			local a = throttle.create(1, 1);
+
+			a.balance = 0;
+			later(1);
+			a:update();
+			assert.same(a, { balance = 1, max = 1, rate = 1, t = now });
+
+			a.balance = 0;
+			for i=1,10 do
+				later(0.1);
+				a:update();
+			end
+			assert(math.abs(a.balance - 1) < 0.0001); -- incremental updates cause rouding errors
+		end);
+	end);
+
+	-- describe("po")
+
+	describe("#poll()", function()
+		it("should only allow successful polls until cost is hit", function()
+			now = 5;
+
+			local a = throttle.create(3, 10);
+			assert.same(a, { balance = 3, max = 3, rate = 0.3, t = 5 });
+
+			assert.is_true(a:poll(1));  -- 3 -> 2
+			assert.same(a, { balance = 2, max = 3, rate = 0.3, t = 5 });
+
+			assert.is_true(a:poll(2));  -- 2 -> 1
+			assert.same(a, { balance = 0, max = 3, rate = 0.3, t = 5 });
+
+			assert.is_false(a:poll(1)); -- MEEP, out of credits!
+			assert.is_false(a:poll(1)); -- MEEP, out of credits!
+			assert.same(a, { balance = 0, max = 3, rate = 0.3, t = 5 });
+		end);
+
+		it("should not allow polls more than the cost", function()
+			now = 0;
+
+			local a = throttle.create(10, 10);
+			assert.same(a, { balance = 10, max = 10, rate = 1, t = 0 });
+
+			assert.is_false(a:poll(11));
+			assert.same(a, { balance = 10, max = 10, rate = 1, t = 0 });
+
+			assert.is_true(a:poll(6));
+			assert.same(a, { balance = 4, max = 10, rate = 1, t = 0 });
+
+			assert.is_false(a:poll(5));
+			assert.same(a, { balance = 4, max = 10, rate = 1, t = 0 });
+
+			-- fractional
+			assert.is_true(a:poll(3.5));
+			assert.same(a, { balance = 0.5, max = 10, rate = 1, t = 0 });
+
+			assert.is_true(a:poll(0.25));
+			assert.same(a, { balance = 0.25, max = 10, rate = 1, t = 0 });
+
+			assert.is_false(a:poll(0.3));
+			assert.same(a, { balance = 0.25, max = 10, rate = 1, t = 0 });
+
+			assert.is_true(a:poll(0.25));
+			assert.same(a, { balance = 0, max = 10, rate = 1, t = 0 });
+
+			assert.is_false(a:poll(0.1));
+			assert.same(a, { balance = 0, max = 10, rate = 1, t = 0 });
+
+			assert.is_true(a:poll(0));
+			assert.same(a, { balance = 0, max = 10, rate = 1, t = 0 });
 		end);
 	end);
 end);
diff --git a/util/throttle.lua b/util/throttle.lua
index 18692a8a..a8191886 100644
--- a/util/throttle.lua
+++ b/util/throttle.lua
@@ -12,7 +12,7 @@ function throttle:update()
 	local newt = gettime();
 	local elapsed = newt - self.t;
 	self.t = newt;
-	local balance = floor(self.rate * elapsed) + self.balance;
+	local balance = (self.rate * elapsed) + self.balance;
 	if balance > self.max then
 		self.balance = self.max;
 	else
@@ -40,7 +40,7 @@ function throttle:poll(cost, split)
 end
 
 local function create(max, period)
-	return setmetatable({ rate = max / period, max = max, t = 0, balance = max }, throttle_mt);
+	return setmetatable({ rate = max / period, max = max, t = gettime(), balance = max }, throttle_mt);
 end
 
 return {