perf(push): batch badge count invalidation per mention chunk (#2178)

fluxer_api/src/api/infrastructure/GatewayService.ts

@@ -32,6 +32,8 @@ import type {
 	GatewayVoiceStateEntry,
 } from './IGatewayService';
 
+const PUSH_BADGE_COUNT_BATCH_SIZE = 100;
+
 const GATEWAY_ERROR_TO_DOMAIN_ERROR: Record<string, () => Error> = {
 	[GatewayRpcMethodErrorCodes.GUILD_NOT_FOUND]: () => new UnknownGuildError(),
 	[GatewayRpcMethodErrorCodes.FORBIDDEN]: () => new MissingPermissionsError(),
@@ -61,6 +63,10 @@ interface InvalidatePushBadgeCountParams {
 	userId: UserID;
 }
 
+interface InvalidatePushBadgeCountsParams {
+	userIds: Array<UserID>;
+}
+
 interface InvalidatePushSubscriptionsParams {
 	userId: UserID;
 }
@@ -578,6 +584,18 @@ export class GatewayService {
 		});
 	}
 
+	async invalidatePushBadgeCounts({userIds}: InvalidatePushBadgeCountsParams): Promise<void> {
+		const batches: Array<Array<UserID>> = [];
+		for (let index = 0; index < userIds.length; index += PUSH_BADGE_COUNT_BATCH_SIZE) {
+			batches.push(userIds.slice(index, index + PUSH_BADGE_COUNT_BATCH_SIZE));
+		}
+		await Promise.all(
+			batches.map((batch) =>
+				this.call('push.invalidate_badge_counts', {user_ids: batch.map((userId) => userId.toString())}),
+			),
+		);
+	}
+
 	async invalidatePushSubscriptions({userId}: InvalidatePushSubscriptionsParams): Promise<void> {
 		await this.call('push.invalidate_subscriptions', {
 			user_id: userId.toString(),

fluxer_api/src/api/infrastructure/IGatewayService.ts

@@ -282,6 +282,8 @@ export abstract class IGatewayService {
 
 	abstract invalidatePushBadgeCount(params: {userId: UserID}): Promise<void>;
 
+	abstract invalidatePushBadgeCounts(params: {userIds: Array<UserID>}): Promise<void>;
+
 	abstract invalidatePushSubscriptions(params: {userId: UserID}): Promise<void>;
 
 	abstract clearPushChannelNotifications(params: {

fluxer_api/src/api/read_state/ReadStateService.test.ts

@@ -0,0 +1,54 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import {describe, expect, it, vi} from 'vitest';
+import {type ChannelID, createChannelID, createMessageID, createUserID, type UserID} from '../BrandedTypes';
+import type {IGatewayService} from '../infrastructure/IGatewayService';
+import type {IReadStateRepository} from './IReadStateRepository';
+import {ReadStateService} from './ReadStateService';
+
+describe('ReadStateService.bulkIncrementMentionCounts', () => {
+	it('invalidates badge counts for touched users in a single bulk call', async () => {
+		const channelId = createChannelID(2n);
+		const messageId = createMessageID(3n);
+		const touched: Array<{userId: UserID; channelId: ChannelID}> = [
+			{userId: createUserID(10n), channelId},
+			{userId: createUserID(11n), channelId},
+			{userId: createUserID(10n), channelId: createChannelID(4n)},
+		];
+		const repository = {
+			bulkIncrementMentionCounts: vi.fn().mockResolvedValue(touched),
+		} as unknown as IReadStateRepository;
+		const invalidatePushBadgeCounts = vi.fn().mockResolvedValue(undefined);
+		const invalidatePushBadgeCount = vi.fn().mockResolvedValue(undefined);
+		const gatewayService = {
+			invalidatePushBadgeCounts,
+			invalidatePushBadgeCount,
+		} as unknown as IGatewayService;
+		const service = new ReadStateService(repository, gatewayService);
+
+		await service.bulkIncrementMentionCounts([
+			{userId: createUserID(10n), channelId, messageId},
+			{userId: createUserID(11n), channelId, messageId},
+			{userId: createUserID(12n), channelId, messageId},
+		]);
+
+		expect(invalidatePushBadgeCount).not.toHaveBeenCalled();
+		expect(invalidatePushBadgeCounts).toHaveBeenCalledTimes(1);
+		expect(invalidatePushBadgeCounts).toHaveBeenCalledWith({userIds: [createUserID(10n), createUserID(11n)]});
+	});
+
+	it('skips the bulk call when no read state was touched', async () => {
+		const repository = {
+			bulkIncrementMentionCounts: vi.fn().mockResolvedValue([]),
+		} as unknown as IReadStateRepository;
+		const invalidatePushBadgeCounts = vi.fn().mockResolvedValue(undefined);
+		const gatewayService = {invalidatePushBadgeCounts} as unknown as IGatewayService;
+		const service = new ReadStateService(repository, gatewayService);
+
+		await service.bulkIncrementMentionCounts([
+			{userId: createUserID(10n), channelId: createChannelID(2n), messageId: createMessageID(3n)},
+		]);
+
+		expect(invalidatePushBadgeCounts).not.toHaveBeenCalled();
+	});
+});

fluxer_api/src/api/read_state/ReadStateService.ts

@@ -171,14 +171,13 @@ export class ReadStateService {
 		try {
 			const appliedUpdates = await this.repository.bulkIncrementMentionCounts(updates);
 			const uniqueUserIds = Array.from(new Set(appliedUpdates.map((update) => update.userId)));
-			await Promise.all(
-				uniqueUserIds.map((userId) =>
-					this.gatewayService.invalidatePushBadgeCount({userId}).catch((error) => {
-						Logger.error({userId: userId.toString(), error}, 'Failed to invalidate push badge count');
-						return null;
-					}),
-				),
-			);
+			if (uniqueUserIds.length === 0) {
+				return;
+			}
+			await this.gatewayService.invalidatePushBadgeCounts({userIds: uniqueUserIds}).catch((error) => {
+				Logger.error({userCount: uniqueUserIds.length, error}, 'Failed to invalidate push badge counts');
+				return null;
+			});
 		} catch (error) {
 			Logger.error({error}, 'Bulk increment mention counts failed');
 			throw error;

fluxer_api/src/api/test/NoopGatewayService.ts

@@ -797,6 +797,8 @@ export class NoopGatewayService extends IGatewayService {
 
 	async invalidatePushBadgeCount(_params: {userId: UserID}): Promise<void> {}
 
+	async invalidatePushBadgeCounts(_params: {userIds: Array<UserID>}): Promise<void> {}
+
 	async invalidatePushSubscriptions(_params: {userId: UserID}): Promise<void> {}
 
 	async clearPushChannelNotifications(_params: {

fluxer_gateway/src/gateway/gateway_rpc_push.erl

@@ -12,6 +12,8 @@ execute_method(<<"push.sync_user_blocked_ids">>, P) ->
     do_sync_user_blocked_ids(P);
 execute_method(<<"push.invalidate_badge_count">>, P) ->
     do_invalidate_badge_count(P);
+execute_method(<<"push.invalidate_badge_counts">>, P) ->
+    do_invalidate_badge_counts(P);
 execute_method(<<"push.clear_channel_notifications">>, P) ->
     do_clear_channel_notifications(P);
 execute_method(<<"push.invalidate_subscriptions">>, P) ->
@@ -43,6 +45,19 @@ do_invalidate_badge_count(#{<<"user_id">> := UBin}) ->
     fanout_local_cache_mutation(push, invalidate_user_badge_count_local, [UserId]),
     true.
 
+-spec do_invalidate_badge_counts(map()) -> true.
+do_invalidate_badge_counts(#{<<"user_ids">> := UserIdsRaw}) ->
+    UserIds = lists:usort(validation:snowflake_list_or_throw(<<"user_ids">>, UserIdsRaw)),
+    gateway_rpc_guild_routing:validate_batch_size(length(UserIds)),
+    fanout_badge_count_invalidation(UserIds).
+
+-spec fanout_badge_count_invalidation([pos_integer()]) -> true.
+fanout_badge_count_invalidation([]) ->
+    true;
+fanout_badge_count_invalidation(UserIds) ->
+    fanout_local_cache_mutation(push, invalidate_user_badge_counts_local, [UserIds]),
+    true.
+
 -spec do_clear_channel_notifications(map()) -> true.
 do_clear_channel_notifications(#{
     <<"user_id">> := UBin,
@@ -79,3 +94,24 @@ fanout_local_cache_mutation(Module, Function, Args) ->
             ),
             ok
     end.
+
+-ifdef(TEST).
+-include_lib("eunit/include/eunit.hrl").
+
+invalidate_badge_counts_rejects_oversized_batch_test() ->
+    UserIds = [integer_to_binary(UserId) || UserId <- lists:seq(1000, 1200)],
+    ?assertError(
+        {gateway_rpc_error, _},
+        execute_method(<<"push.invalidate_badge_counts">>, #{<<"user_ids">> => UserIds})
+    ).
+
+invalidate_badge_counts_rejects_invalid_snowflakes_test() ->
+    ?assertError(
+        {validation, _},
+        execute_method(<<"push.invalidate_badge_counts">>, #{<<"user_ids">> => [<<"nope">>]})
+    ).
+
+invalidate_badge_counts_accepts_an_empty_batch_test() ->
+    ?assert(execute_method(<<"push.invalidate_badge_counts">>, #{<<"user_ids">> => []})).
+
+-endif.

fluxer_gateway/src/push/push.erl

@@ -16,6 +16,7 @@
     invalidate_user_subscriptions_local/1,
     invalidate_user_badge_count/1,
     invalidate_user_badge_count_local/1,
+    invalidate_user_badge_counts_local/1,
     clear_channel_notifications/3
 ]).
 -export([get_cache_stats/0]).
@@ -172,6 +173,15 @@ invalidate_user_badge_count_local(UserId) ->
         push_ets_cache:delete_badge_count(UserId)
     end).
 
+-spec invalidate_user_badge_counts_local(term()) -> ok.
+invalidate_user_badge_counts_local(UserIds) ->
+    case integer_list(UserIds) of
+        {ok, TypedUserIds} ->
+            lists:foreach(fun invalidate_user_badge_count_local/1, TypedUserIds);
+        error ->
+            ok
+    end.
+
 -spec maybe_cast(term(), term()) -> ok.
 maybe_cast(Key, Msg) ->
     case is_push_noop() of
@@ -551,6 +561,25 @@ sync_user_blocked_ids_local_updates_local_cache_test() ->
         ?assertEqual([20, 30], push_ets_cache:get_blocked_ids(10))
     end).
 
+invalidate_user_badge_counts_local_deletes_every_cached_entry_test() ->
+    push_ets_cache:init(),
+    push_ets_cache:put_badge_count(10, 5, 1000),
+    push_ets_cache:put_badge_count(11, 7, 1000),
+    with_registered_push(fun() ->
+        ok = invalidate_user_badge_counts_local([10, 11])
+    end),
+    ?assertEqual(undefined, push_ets_cache:get_badge_count(10)),
+    ?assertEqual(undefined, push_ets_cache:get_badge_count(11)).
+
+invalidate_user_badge_counts_local_ignores_untyped_ids_test() ->
+    push_ets_cache:init(),
+    push_ets_cache:put_badge_count(12, 5, 1000),
+    with_registered_push(fun() ->
+        ok = invalidate_user_badge_counts_local([<<"12">>])
+    end),
+    ?assertEqual({5, 1000}, push_ets_cache:get_badge_count(12)),
+    push_ets_cache:delete_badge_count(12).
+
 invalidate_user_subscriptions_local_deletes_local_cache_test() ->
     push_ets_cache:init(),
     push_ets_cache:put_subscriptions(10, [#{<<"endpoint">> => <<"test">>}]),