chore: add lock method to LockRepository for managing room locks with disposal functionality (#238)

Co-authored-by: Diego Sampaio <chinello@gmail.com>

Author: ggazzo

packages/federation-sdk/src/listeners/staging-area.listener.ts

@@ -13,8 +13,8 @@ export class StagingAreaListener {
 		this.stagingAreaQueue.registerHandler(this.handleQueueItem.bind(this));
 	}
 
-	async handleQueueItem(data: RoomID) {
+	async *handleQueueItem(data: RoomID) {
 		this.logger.debug(`Processing room ${data}`);
-		await this.stagingAreaService.processEventForRoom(data);
+		yield* this.stagingAreaService.processEventForRoom(data);
 	}
 }

packages/federation-sdk/src/queues/staging-area.queue.ts

@@ -1,48 +1,72 @@
 import type { RoomID } from '@rocket.chat/federation-room';
 import 'reflect-metadata';
-import { singleton } from 'tsyringe';
+import { delay, inject, singleton } from 'tsyringe';
 
-type QueueHandler = (roomId: RoomID) => Promise<void>;
+import { LockRepository } from '../repositories/lock.repository';
+import { ConfigService } from '../services/config.service';
+
+type QueueHandler = (roomId: RoomID) => AsyncGenerator<unknown | undefined>;
 
 @singleton()
 export class StagingAreaQueue {
-	private queue: RoomID[] = [];
+	private queue: Set<RoomID> = new Set();
 
-	private handlers: QueueHandler[] = [];
+	private handler: QueueHandler | null = null;
 
 	private processing = false;
 
+	constructor(
+		@inject(delay(() => LockRepository))
+		private readonly lockRepository: LockRepository,
+		private readonly configService: ConfigService,
+	) {}
+
 	enqueue(roomId: RoomID): void {
-		this.queue.push(roomId);
+		this.queue.add(roomId);
 		this.processQueue();
 	}
 
 	registerHandler(handler: QueueHandler): void {
-		this.handlers.push(handler);
+		this.handler = handler;
 	}
 
 	private async processQueue(): Promise<void> {
 		if (this.processing) {
 			return;
 		}
 
+		if (!this.handler) {
+			throw new Error('No handler registered for StagingAreaQueue');
+		}
+
 		this.processing = true;
 
 		try {
-			while (this.queue.length > 0) {
-				const roomId = this.queue.shift() as RoomID;
+			while (this.queue.size > 0) {
+				const [roomId] = this.queue;
 				if (!roomId) continue;
+				this.queue.delete(roomId);
+
+				// eslint-disable-next-line no-await-in-loop, prettier/prettier
+				await using lock = await this.lockRepository.lock(
+					roomId,
+					this.configService.instanceId,
+				);
+
+				if (!lock.success) {
+					continue;
+				}
 
-				for (const handler of this.handlers) {
-					// eslint-disable-next-line no-await-in-loop
-					await handler(roomId);
+				// eslint-disable-next-line no-await-in-loop --- this is valid since this.handler is an async generator
+				for await (const _ of this.handler(roomId)) {
+					await lock.update();
 				}
 			}
 		} finally {
 			this.processing = false;
 
 			// Check if new items were added while processing
-			if (this.queue.length > 0) {
+			if (this.queue.size > 0) {
 				this.processQueue();
 			}
 		}

packages/federation-sdk/src/repositories/lock.repository.ts

@@ -14,6 +14,32 @@ export class LockRepository {
 		this.collection.createIndex({ roomId: 1 }, { unique: true });
 	}
 
+	async lock(
+		roomId: string,
+		instanceId: string,
+	): Promise<{
+		success: boolean;
+		update: () => Promise<void>;
+		[Symbol.asyncDispose]: () => Promise<void>;
+	}> {
+		const lock = await this.getLock(roomId, instanceId);
+		return {
+			success: lock,
+			update: async () => {
+				if (!lock) {
+					return;
+				}
+				return this.updateLockTimestamp(roomId, instanceId);
+			},
+			[Symbol.asyncDispose]: async () => {
+				if (!lock) {
+					return;
+				}
+				return this.releaseLock(roomId, instanceId);
+			},
+		};
+	}
+
 	async getLock(roomId: string, instanceId: string): Promise<boolean> {
 		const timedout = new Date();
 		timedout.setTime(timedout.getTime() - 2 * 60 * 1000); // 2 minutes ago

packages/federation-sdk/src/services/event.service.ts

@@ -26,7 +26,6 @@ import {
 	RoomID,
 	RoomState,
 	RoomVersion,
-	type State,
 	getAuthChain,
 } from '@rocket.chat/federation-room';
 import { delay, inject, singleton } from 'tsyringe';
@@ -39,7 +38,6 @@ import type { StateService } from './state.service';
 import { StagingAreaQueue } from '../queues/staging-area.queue';
 import { EventStagingRepository } from '../repositories/event-staging.repository';
 import { EventRepository } from '../repositories/event.repository';
-import { LockRepository } from '../repositories/lock.repository';
 import { eventSchemas } from '../utils/event-schemas';
 
 export interface AuthEventParams {
@@ -65,8 +63,6 @@ export class EventService {
 		private readonly eventRepository: EventRepository,
 		@inject(delay(() => EventStagingRepository))
 		private readonly eventStagingRepository: EventStagingRepository,
-		@inject(delay(() => LockRepository))
-		private readonly lockRepository: LockRepository,
 	) {}
 
 	async getEventById<T extends PduType, P extends EventStore<PduForType<T>>>(eventId: EventID, type?: T): Promise<P | null> {
@@ -197,16 +193,6 @@ export class EventService {
 					// save the event as staged to be processed
 					await this.eventStagingRepository.create(eventId, origin, event);
 
-					// acquire a lock for processing the event
-					const lock = await this.lockRepository.getLock(roomId, this.configService.instanceId);
-					if (!lock) {
-						this.logger.debug(`Couldn't acquire a lock for room ${roomId}`);
-						continue;
-					}
-
-					// if we have a lock, we can process the event
-					// void this.stagingAreaService.processEventForRoom(roomId);
-
 					// TODO change this to call stagingAreaService directly (line above)
 					this.stagingAreaQueue.enqueue(roomId);
 				}
@@ -637,15 +623,6 @@ export class EventService {
 
 		// not we try to process one room at a time
 		for await (const roomId of rooms) {
-			const lock = await this.lockRepository.getLock(roomId, this.configService.instanceId);
-			if (!lock) {
-				this.logger.debug(`Couldn't acquire a lock for room ${roomId}`);
-				continue;
-			}
-
-			// if we have a lock, we can process the event
-			// void this.stagingAreaService.processEventForRoom(roomId);
-
 			// TODO change this to call stagingAreaService directly (line above)
 			this.stagingAreaQueue.enqueue(roomId);
 

packages/federation-sdk/src/services/staging-area.service.ts

@@ -60,7 +60,7 @@ export class StagingAreaService {
 		return [authEvents, prevEvents];
 	}
 
-	async processEventForRoom(roomId: RoomID) {
+	async *processEventForRoom(roomId: RoomID) {
 		const roomIdToRoomVersion = new Map<string, RoomVersion>();
 		const getRoomVersion = async (roomId: RoomID) => {
 			const version = roomIdToRoomVersion.get(roomId) ?? (await this.stateService.getRoomVersion(roomId));
@@ -92,11 +92,6 @@ export class StagingAreaService {
 
 			this.logger.info({ msg: 'Processing event', eventId: event._id });
 
-			// if we got an event, we need to update the lock's timestamp to avoid it being timed out
-			// and acquired by another instance while we're processing a batch of events for this room
-			// eslint-disable-next-line no-await-in-loop
-			await this.lockRepository.updateLockTimestamp(roomId, this.configService.instanceId);
-
 			try {
 				// eslint-disable-next-line no-await-in-loop
 				const addedMissing = await this.processDependencyStage(event);
@@ -143,10 +138,12 @@ export class StagingAreaService {
 					});
 				}
 			}
-		} while (event);
 
-		// release the lock after processing
-		await this.lockRepository.releaseLock(roomId, this.configService.instanceId);
+			yield event;
+
+			// TODO: what should we do to avoid infinite loops in case the next event is always the same event
+		} while (event);
+		this.logger.debug({ msg: 'No more events to process for room', roomId });
 	}
 
 	private async processDependencyStage(event: EventStagingStore) {