diff teal-src/prosody/util/smqueue.tl @ 14017:1e01b91cf94d

util.smqueue: Backport changes to Teal source
author Kim Alvefur <zash@zash.se>
date Sun, 07 Dec 2025 15:13:09 +0100
parents 9ec961173b1c
children
line wrap: on
line diff
--- a/teal-src/prosody/util/smqueue.tl	Fri Dec 19 15:32:26 2025 +0000
+++ b/teal-src/prosody/util/smqueue.tl	Sun Dec 07 15:13:09 2025 +0100
@@ -6,6 +6,9 @@
 		_queue : queue.queue<T>
 		_head : integer
 		_tail : integer
+		_next_checkpoint : integer
+		type checkpoint = { integer, function(any), any }
+		_checkpoints : { checkpoint }
 
 		enum ack_errors
 			"tail"
@@ -17,6 +20,7 @@
 		resumable : function (smqueue<T>) : boolean
 		resume : function (smqueue<T>)  : queue.queue.iterator, any, integer
 		consume : function (smqueue<T>) : function() : T
+		_call_checkpoints : function(smqueue<T>, integer)
 
 		table : function (smqueue<T>) : { T }
 	end
@@ -31,10 +35,15 @@
 	assert(self._queue:push(v));
 end
 
+-- Call when receiver has acknowledged some stanzas
+-- `h` (sent by the receiver) is the running count of
+-- successfully received stanzas
 function smqueue:ack(h : integer) : { T }, smqueue.ack_errors
 	if h < self._tail then
+		-- h is less than a previous h
 		return nil, "tail"
 	elseif h > self._head then
+		-- h is greater than the number of stanzas we sent
 		return nil, "head"
 	end
 	-- TODO optimize? cache table fields
@@ -43,11 +52,14 @@
 	local expect = self._head - self._tail;
 	while expect < self._queue:count() do
 		local v = self._queue:pop();
-		if not v then
-			return nil, "pop"
-		end
+		if not v then return nil, "pop" end
 		table.insert(acked, v);
 	end
+
+	local c = self._next_checkpoint;
+	if c and h >= c then
+		self:_call_checkpoints(h);
+	end
 	return acked
 end
 
@@ -80,21 +92,54 @@
 	return t
 end
 
+-- Add a 'checkpoint' which will call a function when all the stanzas
+-- currently in the queue have been successfully delivered and acknowledged
+function smqueue:add_checkpoint(cb : function (any), ud : any)
+	table.insert(self._checkpoints, { self._head, cb, ud });
+	if not self._next_checkpoint then
+		self._next_checkpoint = self._head;
+	end
+end
+
+function smqueue:_call_checkpoints(h : integer)
+	local checkpoints = self._checkpoints;
+
+	-- Pop the checkpoint info and update the next checkpoint info
+	local c = table.remove(checkpoints, 1);
+	self._next_checkpoint = checkpoints[1] and checkpoints[1][1] or nil;
+
+	-- Call the callback
+	local cb, ud = c[2], c[3];
+	cb(ud);
+
+	-- In case we passed over multiple checkpoints, call them also
+	if h >= self._next_checkpoint then
+		return self:_call_checkpoints(h)
+	end
+end
+
 local function freeze(q : smqueue<any>) : { string:integer }
 	return { head = q._head, tail = q._tail }
 end
 
-local queue_mt = {
+local smqueue_mt = {
 	--
 	__name = "smqueue";
 	__index = smqueue;
 	__len = smqueue.count_unacked;
+	-- Return a simple object for serialization
+	-- This is forbidden by Teal. Thanks Teal!
 	__freeze = freeze;
 }
 
 function lib.new<T>(size : integer) : smqueue<T>
 	assert(size>0);
-	return setmetatable({ _head = 0; _tail = 0; _queue = queue.new(size, true) }, queue_mt)
+	return setmetatable({
+		_head = 0;
+		_tail = 0;
+		_queue = queue.new(size, true);
+		_checkpoints = {};
+	}, smqueue_mt)
 end
 
 return lib