refactor: state calculation (#252)

Co-authored-by: Guilherme Gazzo <[email protected]>

Author: debdutdeb

.github/workflows/ci.yml

@@ -1,14 +1,15 @@
-
 on:
+  release:
+    types: [published]
   pull_request:
-    branches: '**'
+    branches: "**"
     paths-ignore:
-      - '**.md'
+      - "**.md"
   push:
     branches:
       - main
     paths-ignore:
-      - '**.md'
+      - "**.md"
 
 name: my-workflow
 jobs:
@@ -26,15 +27,14 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-turbo-
       - uses: oven-sh/setup-bun@v2
-
       - run: bun install
       - run: bun run build
       - run: bun lint:ci
       - run: bun tsc --noEmit
-      - run: bun test:coverage
-
+      - uses: supercharge/[email protected]
+      - run: LOG_LEVEL=debug RUN_MONGO_TESTS=1 bun test packages/federation-sdk/src/services/state.service.spec.ts
+      - run: LOG_LEVEL=debug bun test:coverage
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v5
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
-

packages/core/src/models/event.model.ts

@@ -1,4 +1,9 @@
-import { EventID, Pdu, PduForType } from '@rocket.chat/federation-room';
+import type {
+	EventID,
+	Pdu,
+	RejectCode,
+	StateID,
+} from '@rocket.chat/federation-room';
 import type { EventBase as CoreEventBase } from '../events/eventBase';
 
 // TODO: use room package
@@ -19,9 +24,15 @@ interface PersistentEventBase<E = Pdu> {
 
 // TODO: Merge with StagedEvent from event.service.ts
 export interface EventStore<E = Pdu> extends PersistentEventBase<E> {
-	stateId: string;
+	stateId: StateID;
 	// for prev_events
-	nextEventId: string;
+	nextEventId: EventID | '';
+
+	rejectCode?: RejectCode;
+	rejectDetail?: {
+		reason: string;
+		rejectedBy?: EventID;
+	};
 }
 
 export interface EventStagingStore extends PersistentEventBase {

packages/federation-sdk/src/container.ts

@@ -5,7 +5,7 @@ import type {
 	EventStagingStore,
 	EventStore,
 } from '@rocket.chat/federation-core';
-import type { Collection, WithId } from 'mongodb';
+import type { Collection } from 'mongodb';
 import { container } from 'tsyringe';
 
 import type { HomeserverEventSignatures } from './index';
@@ -17,7 +17,10 @@ import { Key, KeyRepository } from './repositories/key.repository';
 import { Lock, LockRepository } from './repositories/lock.repository';
 import { Room, RoomRepository } from './repositories/room.repository';
 import { Server, ServerRepository } from './repositories/server.repository';
-import { StateRepository, StateStore } from './repositories/state.repository';
+import {
+	StateGraphRepository,
+	StateGraphStore,
+} from './repositories/state-graph.repository';
 import { Upload, UploadRepository } from './repositories/upload.repository';
 import { ConfigService } from './services/config.service';
 import { DatabaseConnectionService } from './services/database-connection.service';
@@ -82,10 +85,6 @@ export async function createFederationContainer(
 		useValue: db.collection<Room>('rocketchat_federation_rooms'),
 	});
 
-	container.register<Collection<WithId<StateStore>>>('StateCollection', {
-		useValue: db.collection<WithId<StateStore>>('rocketchat_federation_states'),
-	});
-
 	container.register<Collection<Server>>('ServerCollection', {
 		useValue: db.collection<Server>('rocketchat_federation_servers'),
 	});
@@ -94,14 +93,20 @@ export async function createFederationContainer(
 		useValue: db.collection<Upload>('rocketchat_uploads'),
 	});
 
+	container.register<Collection<StateGraphStore>>('StateGraphCollection', {
+		useValue: db.collection<StateGraphStore>(
+			'rocketchat_federation_state_graphs',
+		),
+	});
+
 	container.registerSingleton(EventRepository);
 	container.registerSingleton(EventStagingRepository);
 	container.registerSingleton(KeyRepository);
 	container.registerSingleton(LockRepository);
 	container.registerSingleton(RoomRepository);
-	container.registerSingleton(StateRepository);
 	container.registerSingleton(ServerRepository);
 	container.registerSingleton(UploadRepository);
+	container.registerSingleton(StateGraphRepository);
 
 	container.registerSingleton(FederationRequestService);
 	container.registerSingleton(FederationService);

packages/federation-sdk/src/index.ts

@@ -4,6 +4,7 @@ import { container } from 'tsyringe';
 import { ConfigService } from './services/config.service';
 import { EduService } from './services/edu.service';
 import { EventAuthorizationService } from './services/event-authorization.service';
+import { EventEmitterService } from './services/event-emitter.service';
 import { EventService } from './services/event.service';
 import { FederationRequestService } from './services/federation-request.service';
 import { InviteService } from './services/invite.service';
@@ -94,7 +95,6 @@ export { EventRepository } from './repositories/event.repository';
 export { RoomRepository } from './repositories/room.repository';
 export { ServerRepository } from './repositories/server.repository';
 export { KeyRepository } from './repositories/key.repository';
-export { StateRepository } from './repositories/state.repository';
 
 export interface HomeserverServices {
 	room: RoomService;
@@ -111,6 +111,7 @@ export interface HomeserverServices {
 	media: MediaService;
 	request: FederationRequestService;
 	federationAuth: EventAuthorizationService;
+	emitter: EventEmitterService;
 }
 
 export type HomeserverEventSignatures = {
@@ -292,6 +293,7 @@ export function getAllServices(): HomeserverServices {
 		media: container.resolve(MediaService),
 		request: container.resolve(FederationRequestService),
 		federationAuth: container.resolve(EventAuthorizationService),
+		emitter: container.resolve(EventEmitterService),
 	};
 }
 

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

@@ -5,9 +5,16 @@ import {
 	Pdu,
 	PduForType,
 	PduType,
+	RejectCode,
+	StateID,
 } from '@rocket.chat/federation-room';
-import type { Collection, FindCursor, WithId } from 'mongodb';
-import { MongoError } from 'mongodb';
+import type {
+	Collection,
+	FindCursor,
+	FindOptions,
+	UpdateResult,
+	WithId,
+} from 'mongodb';
 import { inject, singleton } from 'tsyringe';
 
 @singleton()
@@ -110,15 +117,6 @@ export class EventRepository {
 		);
 	}
 
-	async create(
-		origin: string,
-		event: Pdu,
-		eventId: EventID,
-		stateId = '',
-	): Promise<string | undefined> {
-		return this.persistEvent(origin, event, eventId, stateId);
-	}
-
 	async redactEvent(eventId: EventID, redactedEvent: Pdu): Promise<void> {
 		await this.collection.updateOne(
 			{ _id: eventId },
@@ -168,7 +166,7 @@ export class EventRepository {
 			{
 				'event.room_id': roomId,
 				'event.origin_server_ts': { $lt: timestamp }, // events before passed timestamp
-				stateId: { $ne: '' },
+				stateId: { $ne: '' as StateID },
 			},
 			{
 				sort: {
@@ -192,58 +190,25 @@ export class EventRepository {
 			});
 	}
 
-	async updateStateId(eventId: EventID, stateId: string): Promise<void> {
+	async updateStateId(eventId: EventID, stateId: StateID): Promise<void> {
 		await this.collection.updateOne({ _id: eventId }, { $set: { stateId } });
 	}
 
 	// finds events not yet referenced by other events
 	// more on the respective adr
-	async findPrevEvents(roomId: string) {
+	async findLatestEvents(roomId: string) {
 		return this.collection
-			.find({
-				nextEventId: '',
-				'event.room_id': roomId,
-				_id: { $ne: '' as EventID },
-			})
+			.find(
+				{
+					nextEventId: '',
+					'event.room_id': roomId,
+					rejectCode: { $exists: false },
+				},
+				{ sort: { 'event.depth': 1, createdAt: 1 } },
+			)
 			.toArray();
 	}
 
-	private async persistEvent(
-		origin: string,
-		event: Pdu,
-		eventId: EventID,
-		stateId: string,
-	) {
-		try {
-			await this.collection.insertOne({
-				_id: eventId,
-				origin,
-				event: event,
-				stateId: stateId,
-				createdAt: new Date(),
-				nextEventId: '', // new events are not expected to have forward edges
-			});
-		} catch (e) {
-			if (e instanceof MongoError) {
-				if (e.code === 11000) {
-					// duplicate key error
-					// this is expected, if the same intentional event is attempted to be persisted again
-					return;
-				}
-			}
-
-			throw e;
-		}
-
-		// this must happen later to as to avoid finding 0 prev_events on a parallel request
-		await this.collection.updateMany(
-			{ _id: { $in: event.prev_events } },
-			{ $set: { nextEventId: eventId } },
-		);
-
-		return eventId;
-	}
-
 	findMembershipEventsFromDirectMessageRooms(
 		users: string[],
 	): FindCursor<EventStore<PduForType<'m.room.member'>>> {
@@ -266,9 +231,23 @@ export class EventRepository {
 	findByIds<T extends PduType>(
 		eventIds: EventID[],
 	): FindCursor<WithId<EventStore<PduForType<T>>>> {
-		return this.collection.find({
-			_id: { $in: eventIds },
-		}) as FindCursor<WithId<EventStore<PduForType<T>>>>;
+		return this.collection.find(
+			{
+				_id: { $in: eventIds },
+			},
+			{ sort: { 'event.depth': -1, createdAt: -1 } },
+		) as FindCursor<WithId<EventStore<PduForType<T>>>>;
+	}
+
+	findByIdsOrderedDescending<T extends PduType>(
+		eventIds: EventID[],
+	): FindCursor<WithId<EventStore<PduForType<T>>>> {
+		return this.collection.find(
+			{
+				_id: { $in: eventIds },
+			},
+			{ sort: { 'event.depth': -1, createdAt: -1 } },
+		) as FindCursor<WithId<EventStore<PduForType<T>>>>;
 	}
 
 	findByRoomIdAndTypes(
@@ -422,4 +401,88 @@ export class EventRepository {
 			.sort({ 'event.depth': -1, 'event.origin_server_ts': -1 })
 			.limit(limit);
 	}
+
+	// new ones
+	// -------------------
+
+	insertOrUpdateEventWithStateId(
+		eventId: EventID,
+		event: Pdu,
+		stateId: StateID,
+	): Promise<UpdateResult> {
+		return this.collection.updateOne(
+			{ _id: eventId },
+			{
+				$setOnInsert: {
+					event,
+					nextEventId: '',
+					createdAt: new Date(),
+				},
+				$set: {
+					stateId,
+				},
+			},
+			{ upsert: true },
+		);
+	}
+
+	async updateNextEventReferences(
+		newEventId: EventID,
+		previousEventIds: EventID[],
+	): Promise<UpdateResult> {
+		return this.collection.updateMany(
+			{ _id: { $in: previousEventIds }, nextEventId: '' as EventID },
+			{ $set: { nextEventId: newEventId } },
+		);
+	}
+
+	async findStateIdByEventId(eventId: EventID): Promise<StateID | undefined> {
+		const result = await this.collection.findOne<Pick<EventStore, 'stateId'>>(
+			{ _id: eventId },
+			{ projection: { stateId: 1 } },
+		);
+
+		return result?.stateId;
+	}
+
+	async rejectEvent(
+		eventId: EventID,
+		event: Pdu,
+		stateId: StateID,
+		code: RejectCode,
+		reason: string,
+		rejectedBy?: EventID,
+	): Promise<UpdateResult> {
+		return this.collection.updateOne(
+			{
+				_id: eventId,
+			},
+			{
+				$setOnInsert: {
+					event,
+					stateId,
+					nextEventId: '',
+				},
+				$set: {
+					rejectCode: code,
+					rejectDetail: {
+						reason,
+						rejectedBy,
+					},
+				},
+			},
+			{ upsert: true },
+		);
+	}
+
+	findStateIdsByEventIds(eventIds: EventID[]) {
+		return this.collection.find<Pick<EventStore, 'stateId'>>(
+			{ _id: { $in: eventIds }, stateId: { $ne: '' as StateID } },
+			{ projection: { stateId: 1 } },
+		);
+	}
+
+	findByType(type: PduType) {
+		return this.collection.find({ 'event.type': type });
+	}
 }