fix: back off and preserve status when remote media download fails

packages/federation-sdk/src/services/media.service.spec.ts

@@ -0,0 +1,110 @@
+import { describe, expect, it, mock } from 'bun:test';
+
+import type { ConfigService } from './config.service';
+import { FederationRequestError, type FederationRequestService } from './federation-request.service';
+import { MediaDownloadError, MediaService } from './media.service';
+
+const requestError = (status: number) =>
+	new FederationRequestError(
+		{
+			ok: false,
+			status,
+			headers: {},
+			buffer: () => Promise.reject(),
+			json: () => Promise.reject(),
+			text: () => Promise.reject(),
+			multipart: () => Promise.reject(),
+			body: () => Promise.reject(),
+		},
+		'',
+	);
+
+const buildService = (requestBinaryData: unknown) =>
+	new MediaService({} as ConfigService, { requestBinaryData } as unknown as FederationRequestService);
+
+describe('MediaService.downloadFromRemoteServer', () => {
+	it('returns the content from the first endpoint that answers', async () => {
+		const content = Buffer.from('file-bytes');
+		const requestBinaryData = mock(async () => ({ content }));
+
+		const result = await buildService(requestBinaryData).downloadFromRemoteServer('remote.example', 'abc');
+
+		expect(result).toBe(content);
+		expect(requestBinaryData).toHaveBeenCalledTimes(1);
+	});
+
+	it('falls through to the legacy endpoints before giving up', async () => {
+		const content = Buffer.from('file-bytes');
+		const requestBinaryData = mock(async () => {
+			if (requestBinaryData.mock.calls.length < 3) {
+				throw requestError(404);
+			}
+			return { content };
+		});
+
+		const result = await buildService(requestBinaryData).downloadFromRemoteServer('remote.example', 'abc');
+
+		expect(result).toBe(content);
+		expect(requestBinaryData).toHaveBeenCalledTimes(3);
+	});
+
+	it('reports the status of every attempt when all endpoints fail', async () => {
+		const requestBinaryData = mock(async () => {
+			throw requestError(404);
+		});
+
+		const downloadPromise = buildService(requestBinaryData).downloadFromRemoteServer('remote.example', 'abc');
+
+		expect(downloadPromise).rejects.toBeInstanceOf(MediaDownloadError);
+		expect(downloadPromise).rejects.toMatchObject({
+			statuses: [404, 404, 404],
+			status: 404,
+			serverName: 'remote.example',
+			mediaId: 'abc',
+		});
+	});
+
+	it('treats a 404 as retryable, because a large upload may not be committed yet', async () => {
+		const requestBinaryData = mock(async () => {
+			throw requestError(404);
+		});
+
+		const downloadPromise = buildService(requestBinaryData).downloadFromRemoteServer('remote.example', 'abc');
+
+		expect(downloadPromise).rejects.toMatchObject({ retryable: true });
+	});
+
+	it.each([401, 403, 410])('treats a %d refusal or removal as not retryable', async (status) => {
+		const requestBinaryData = mock(async () => {
+			throw requestError(status);
+		});
+
+		const downloadPromise = buildService(requestBinaryData).downloadFromRemoteServer('remote.example', 'abc');
+
+		expect(downloadPromise).rejects.toMatchObject({ retryable: false });
+	});
+
+	it('treats a server error as retryable', async () => {
+		const requestBinaryData = mock(async () => {
+			throw requestError(502);
+		});
+
+		const downloadPromise = buildService(requestBinaryData).downloadFromRemoteServer('remote.example', 'abc');
+
+		expect(downloadPromise).rejects.toMatchObject({ retryable: true });
+	});
+
+	it('treats a failure with no response as retryable', async () => {
+		const requestBinaryData = mock(async () => {
+			throw new Error('socket hang up');
+		});
+
+		const downloadPromise = buildService(requestBinaryData).downloadFromRemoteServer('remote.example', 'abc');
+
+		expect(downloadPromise).rejects.toMatchObject({
+			statuses: [undefined, undefined, undefined],
+			status: undefined,
+			retryable: true,
+		});
+	});
+});

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

@@ -2,7 +2,34 @@ import { createLogger } from '@rocket.chat/federation-core';
 import { singleton } from 'tsyringe';
 
 import { ConfigService } from './config.service';
-import { FederationRequestService } from './federation-request.service';
+import { FederationRequestError, FederationRequestService } from './federation-request.service';
+
+export class MediaDownloadError extends Error {
+	readonly name = 'MediaDownloadError';
+
+	constructor(readonly serverName: string, readonly mediaId: string, readonly statuses: (number | undefined)[]) {
+		super(`Failed to download media ${mediaId} from ${serverName} (statuses: ${statuses.map((s) => s ?? 'network').join(', ')})`);
+	}
+
+	get status(): number | undefined {
+		return this.statuses[this.statuses.length - 1];
+	}
+
+	get retryable(): boolean {
+		return this.statuses.some((status) => {
+			if (status === undefined) {
+				// no response at all: transport problem, worth another attempt
+				return true;
+			}
+
+			if (status === 404 || status === 408 || status === 429) {
+				return true;
+			}
+
+			return status >= 500;
+		});
+	}
+}
 
 @singleton()
 export class MediaService {
@@ -17,17 +44,21 @@ export class MediaService {
 			`/_matrix/media/r0/download/${serverName}/${mediaId}`,
 		];
 
+		const statuses: (number | undefined)[] = [];
+
 		for await (const endpoint of endpoints) {
 			try {
 				// TODO: Stream remote file downloads instead of buffering the entire file in memory.
 				const response = await this.federationRequest.requestBinaryData('GET', serverName, endpoint);
 
 				return response.content;
 			} catch (err) {
+				const status = err instanceof FederationRequestError ? err.response.status : undefined;
+				statuses.push(status);
 				this.logger.debug(`Endpoint ${endpoint} failed: ${err instanceof Error ? err.message : String(err)}`);
 			}
 		}
 
-		throw new Error(`Failed to download media ${mediaId} from ${serverName}`);
+		throw new MediaDownloadError(serverName, mediaId, statuses);
 	}
 }

packages/federation-sdk/src/services/staging-area-retry.spec.ts

@@ -0,0 +1,34 @@
+import { describe, expect, it } from 'bun:test';
+
+import { retryDelayForAttempt } from './staging-area.service';
+
+describe('retryDelayForAttempt', () => {
+	it('does not delay a first attempt', () => {
+		expect(retryDelayForAttempt(0)).toBe(0);
+	});
+
+	it('grows exponentially from the base delay', () => {
+		expect(retryDelayForAttempt(1)).toBe(500);
+		expect(retryDelayForAttempt(2)).toBe(1000);
+		expect(retryDelayForAttempt(3)).toBe(2000);
+		expect(retryDelayForAttempt(4)).toBe(4000);
+	});
+
+	it('caps the delay so a room pass cannot outlive its watchdog', () => {
+		expect(retryDelayForAttempt(5)).toBe(5000);
+		expect(retryDelayForAttempt(10)).toBe(5000);
+		expect(retryDelayForAttempt(100)).toBe(5000);
+	});
+
+	it('spreads a full retry budget across seconds, not milliseconds', () => {
+		const total = Array.from({ length: 10 }, (_, i) => retryDelayForAttempt(i + 1)).reduce((a, b) => a + b, 0);
+
+		expect(total).toBeGreaterThan(30_000);
+	});
+
+	it('never delays a single attempt longer than the cap', () => {
+		for (let attempt = 0; attempt <= 20; attempt++) {
+			expect(retryDelayForAttempt(attempt)).toBeLessThanOrEqual(5000);
+		}
+	});
+});

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

@@ -0,0 +1,166 @@
+import { describe, expect, it, mock } from 'bun:test';
+
+import type { EventStagingStore, EventStore } from '@rocket.chat/federation-core';
+import type { EventID, Pdu, PersistentEventBase, RoomID, RoomVersion, UserID } from '@rocket.chat/federation-room';
+
+import type { ConfigService } from './config.service';
+import type { EventAuthorizationService } from './event-authorization.service';
+import type { EventEmitterService } from './event-emitter.service';
+import type { EventService } from './event.service';
+import type { FederationService } from './federation.service';
+import type { MissingEventService } from './missing-event.service';
+import { StagingAreaService } from './staging-area.service';
+import type { StateService } from './state.service';
+import type { LockRepository } from '../repositories/lock.repository';
+
+const ROOM_ID = '!room:example.com' as RoomID;
+
+const stagedEvent = (id: string, got: number): EventStagingStore => ({
+	_id: id as EventID,
+	got,
+	roomId: ROOM_ID,
+	origin: 'remote.example',
+	from: 'transaction',
+	createdAt: new Date(0),
+	event: {
+		type: 'm.room.message',
+		room_id: ROOM_ID,
+		sender: '@alice:remote.example' as UserID,
+		content: { msgtype: 'm.text', body: 'hi' },
+		depth: 5,
+		prev_events: [],
+		auth_events: [],
+		origin_server_ts: 1,
+		hashes: { sha256: 'x' },
+		signatures: {},
+	},
+});
+
+type StagedEventService = Pick<
+	EventService,
+	'getLeastDepthEventForRoom' | 'markEventAsUnstaged' | 'authorizeAndPersist' | 'checkIfEventsExists' | 'getLastEventForRoom'
+>;
+
+const buildService = (queue: (EventStagingStore | null)[], overrides: Partial<StagedEventService> = {}) => {
+	const getLeastDepthEventForRoom = mock(async (_roomId: string) => queue.shift() ?? null);
+	const markEventAsUnstaged = mock(async (_event: EventStagingStore) => undefined);
+	const authorizeAndPersist = mock(async (_event: PersistentEventBase): Promise<void> => {
+		throw new Error('Failed to download media abc from remote.example');
+	});
+
+	const eventService = {
+		getLeastDepthEventForRoom,
+		markEventAsUnstaged,
+		authorizeAndPersist,
+		checkIfEventsExists: mock(async (_eventIds: EventID[]) => ({
+			missing: [] as EventID[],
+			found: [] as EventID[],
+		})),
+		getLastEventForRoom: mock(async (_roomId: string): Promise<EventStore | null> => null),
+		...overrides,
+	} satisfies StagedEventService as unknown as EventService;
+
+	const service = new StagingAreaService(
+		{} as ConfigService,
+		eventService,
+		{
+			fetchMissingEvent: mock(async (): Promise<boolean> => false),
+		} as unknown as MissingEventService,
+		{} as EventAuthorizationService,
+		{} as EventEmitterService,
+		{
+			getRoomVersion: mock(async (_roomId: RoomID): Promise<RoomVersion> => '10'),
+		} as unknown as StateService,
+		{
+			getMissingEvents: mock(async (): Promise<{ events: Pdu[] }> => ({ events: [] })),
+		} as unknown as FederationService,
+		{} as LockRepository,
+	);
+
+	return {
+		service,
+		getLeastDepthEventForRoom,
+		markEventAsUnstaged,
+		authorizeAndPersist,
+	};
+};
+
+const drain = async (generator: ReturnType<StagingAreaService['processEventForRoom']>) => {
+	const seen = [];
+	for await (const item of generator) {
+		seen.push(item);
+	}
+	return seen;
+};
+
+describe('StagingAreaService.processEventForRoom', () => {
+	it('leaves a failing event staged instead of consuming it', async () => {
+		const event = stagedEvent('$a', 0);
+		const { service, markEventAsUnstaged, authorizeAndPersist } = buildService([event, event]);
+
+		await drain(service.processEventForRoom(ROOM_ID));
+
+		expect(authorizeAndPersist).toHaveBeenCalledTimes(1);
+		expect(markEventAsUnstaged, 'the event is only unstaged on success or once the retry budget is spent').not.toHaveBeenCalled();
+	});
+
+	it('attempts an event at most once per pass', async () => {
+		const event = stagedEvent('$a', 0);
+		const { service, authorizeAndPersist } = buildService([event, event, event, event, event]);
+
+		await drain(service.processEventForRoom(ROOM_ID));
+
+		expect(authorizeAndPersist).toHaveBeenCalledTimes(1);
+	});
+
+	it('still processes other staged events in the same pass', async () => {
+		const first = stagedEvent('$a', 0);
+		const second = stagedEvent('$b', 0);
+		const { service, authorizeAndPersist } = buildService([first, second, first]);
+
+		await drain(service.processEventForRoom(ROOM_ID));
+
+		expect(authorizeAndPersist).toHaveBeenCalledTimes(2);
+	});
+
+	it('waits before retrying an event that already failed', async () => {
+		const retried = stagedEvent('$a', 1);
+		const { service } = buildService([retried, retried]);
+
+		const startedAt = Date.now();
+		await drain(service.processEventForRoom(ROOM_ID));
+
+		expect(Date.now() - startedAt).toBeGreaterThanOrEqual(400);
+	});
+
+	it('does not wait before a first attempt', async () => {
+		const fresh = stagedEvent('$a', 0);
+		const { service } = buildService([fresh, fresh]);
+
+		const startedAt = Date.now();
+		await drain(service.processEventForRoom(ROOM_ID));
+
+		expect(Date.now() - startedAt).toBeLessThan(300);
+	});
+
+	it('drops an event once the retry budget is exhausted', async () => {
+		const exhausted = stagedEvent('$a', 999);
+		const { service, markEventAsUnstaged, authorizeAndPersist } = buildService([exhausted]);
+
+		await drain(service.processEventForRoom(ROOM_ID));
+
+		expect(markEventAsUnstaged).toHaveBeenCalledTimes(1);
+		expect(authorizeAndPersist).not.toHaveBeenCalled();
+	});
+
+	it('unstages an event that processes successfully', async () => {
+		const event = stagedEvent('$a', 0);
+		const { service, markEventAsUnstaged } = buildService([event, null], {
+			authorizeAndPersist: mock(async () => undefined),
+		});
+
+		await drain(service.processEventForRoom(ROOM_ID));
+
+		expect(markEventAsUnstaged).toHaveBeenCalledTimes(1);
+	});
+});