From 3238815b3051574fc37d5e230a461e37d300e540 Mon Sep 17 00:00:00 2001 From: Kagami Sascha Rosylight Date: Sun, 7 Jan 2024 22:02:15 +0100 Subject: [PATCH] refactor: make sure promises are settled before app shutdown --- packages/backend/src/GlobalModule.ts | 17 +++++--------- .../backend/src/core/NoteCreateService.ts | 3 ++- packages/backend/src/core/NoteReadService.ts | 9 ++++---- .../backend/src/core/NotificationService.ts | 14 ++++++++++- packages/backend/src/core/QueueModule.ts | 13 ++++------- packages/backend/src/core/ReactionService.ts | 5 ++-- .../activitypub/ApDeliverManagerService.ts | 2 +- packages/backend/src/misc/promise-tracker.ts | 23 +++++++++++++++++++ .../server/api/endpoints/antennas/notes.ts | 3 ++- 9 files changed, 59 insertions(+), 30 deletions(-) create mode 100644 packages/backend/src/misc/promise-tracker.ts diff --git a/packages/backend/src/GlobalModule.ts b/packages/backend/src/GlobalModule.ts index 3e9d19f825..c83845b94c 100644 --- a/packages/backend/src/GlobalModule.ts +++ b/packages/backend/src/GlobalModule.ts @@ -3,7 +3,6 @@ * SPDX-License-Identifier: AGPL-3.0-only */ -import { setTimeout } from 'node:timers/promises'; import { Global, Inject, Module } from '@nestjs/common'; import * as Redis from 'ioredis'; import { DataSource } from 'typeorm'; @@ -12,6 +11,7 @@ import { DI } from './di-symbols.js'; import { Config, loadConfig } from './config.js'; import { createPostgresDataSource } from './postgres.js'; import { RepositoryModule } from './models/RepositoryModule.js'; +import { allSettled } from './misc/promise-tracker.js'; import type { Provider, OnApplicationShutdown } from '@nestjs/common'; const $config: Provider = { @@ -33,7 +33,7 @@ const $meilisearch: Provider = { useFactory: (config: Config) => { if (config.meilisearch) { return new MeiliSearch({ - host: `${config.meilisearch.ssl ? 'https' : 'http' }://${config.meilisearch.host}:${config.meilisearch.port}`, + host: `${config.meilisearch.ssl ? 'https' : 'http'}://${config.meilisearch.host}:${config.meilisearch.port}`, apiKey: config.meilisearch.apiKey, }); } else { @@ -91,17 +91,12 @@ export class GlobalModule implements OnApplicationShutdown { @Inject(DI.redisForPub) private redisForPub: Redis.Redis, @Inject(DI.redisForSub) private redisForSub: Redis.Redis, @Inject(DI.redisForTimelines) private redisForTimelines: Redis.Redis, - ) {} + ) { } public async dispose(): Promise { - if (process.env.NODE_ENV === 'test') { - // XXX: - // Shutting down the existing connections causes errors on Jest as - // Misskey has asynchronous postgres/redis connections that are not - // awaited. - // Let's wait for some random time for them to finish. - await setTimeout(5000); - } + // Wait for all potential DB queries + await allSettled(); + // And then disconnect from DB await Promise.all([ this.db.destroy(), this.redisClient.disconnect(), diff --git a/packages/backend/src/core/NoteCreateService.ts b/packages/backend/src/core/NoteCreateService.ts index ed8d51df16..97fb80ab39 100644 --- a/packages/backend/src/core/NoteCreateService.ts +++ b/packages/backend/src/core/NoteCreateService.ts @@ -58,6 +58,7 @@ import { FanoutTimelineService } from '@/core/FanoutTimelineService.js'; import { UtilityService } from '@/core/UtilityService.js'; import { UserBlockingService } from '@/core/UserBlockingService.js'; import { isReply } from '@/misc/is-reply.js'; +import { trackPromise } from '@/misc/promise-tracker.js'; type NotificationType = 'reply' | 'renote' | 'quote' | 'mention'; @@ -676,7 +677,7 @@ export class NoteCreateService implements OnApplicationShutdown { this.relayService.deliverToRelays(user, noteActivity); } - dm.execute(); + trackPromise(dm.execute()); })(); } //#endregion diff --git a/packages/backend/src/core/NoteReadService.ts b/packages/backend/src/core/NoteReadService.ts index 03c1735e04..c73cf76592 100644 --- a/packages/backend/src/core/NoteReadService.ts +++ b/packages/backend/src/core/NoteReadService.ts @@ -14,6 +14,7 @@ import { IdService } from '@/core/IdService.js'; import { GlobalEventService } from '@/core/GlobalEventService.js'; import type { NoteUnreadsRepository, MutingsRepository, NoteThreadMutingsRepository } from '@/models/_.js'; import { bindThis } from '@/decorators.js'; +import { trackPromise } from '@/misc/promise-tracker.js'; @Injectable() export class NoteReadService implements OnApplicationShutdown { @@ -107,7 +108,7 @@ export class NoteReadService implements OnApplicationShutdown { // TODO: ↓まとめてクエリしたい - this.noteUnreadsRepository.countBy({ + trackPromise(this.noteUnreadsRepository.countBy({ userId: userId, isMentioned: true, }).then(mentionsCount => { @@ -115,9 +116,9 @@ export class NoteReadService implements OnApplicationShutdown { // 全て既読になったイベントを発行 this.globalEventService.publishMainStream(userId, 'readAllUnreadMentions'); } - }); + })); - this.noteUnreadsRepository.countBy({ + trackPromise(this.noteUnreadsRepository.countBy({ userId: userId, isSpecified: true, }).then(specifiedCount => { @@ -125,7 +126,7 @@ export class NoteReadService implements OnApplicationShutdown { // 全て既読になったイベントを発行 this.globalEventService.publishMainStream(userId, 'readAllUnreadSpecifiedNotes'); } - }); + })); } } diff --git a/packages/backend/src/core/NotificationService.ts b/packages/backend/src/core/NotificationService.ts index ad7be83e5b..765fcae063 100644 --- a/packages/backend/src/core/NotificationService.ts +++ b/packages/backend/src/core/NotificationService.ts @@ -20,6 +20,7 @@ import { CacheService } from '@/core/CacheService.js'; import type { Config } from '@/config.js'; import { UserListService } from '@/core/UserListService.js'; import type { FilterUnionByProperty } from '@/types.js'; +import { trackPromise } from '@/misc/promise-tracker.js'; @Injectable() export class NotificationService implements OnApplicationShutdown { @@ -74,7 +75,18 @@ export class NotificationService implements OnApplicationShutdown { } @bindThis - public async createNotification( + public createNotification( + notifieeId: MiUser['id'], + type: T, + data: Omit, 'type' | 'id' | 'createdAt' | 'notifierId'>, + notifierId?: MiUser['id'] | null, + ) { + trackPromise( + this.#createNotificationInternal(notifieeId, type, data, notifierId), + ); + } + + async #createNotificationInternal( notifieeId: MiUser['id'], type: T, data: Omit, 'type' | 'id' | 'createdAt' | 'notifierId'>, diff --git a/packages/backend/src/core/QueueModule.ts b/packages/backend/src/core/QueueModule.ts index 4444dc9787..20a53ff282 100644 --- a/packages/backend/src/core/QueueModule.ts +++ b/packages/backend/src/core/QueueModule.ts @@ -3,12 +3,12 @@ * SPDX-License-Identifier: AGPL-3.0-only */ -import { setTimeout } from 'node:timers/promises'; import { Inject, Module, OnApplicationShutdown } from '@nestjs/common'; import * as Bull from 'bullmq'; import { DI } from '@/di-symbols.js'; import type { Config } from '@/config.js'; import { QUEUE, baseQueueOptions } from '@/queue/const.js'; +import { allSettled } from '@/misc/promise-tracker.js'; import type { Provider } from '@nestjs/common'; import type { DeliverJobData, InboxJobData, EndedPollNotificationJobData, WebhookDeliverJobData, RelationshipJobData } from '../queue/types.js'; @@ -106,14 +106,9 @@ export class QueueModule implements OnApplicationShutdown { ) {} public async dispose(): Promise { - if (process.env.NODE_ENV === 'test') { - // XXX: - // Shutting down the existing connections causes errors on Jest as - // Misskey has asynchronous postgres/redis connections that are not - // awaited. - // Let's wait for some random time for them to finish. - await setTimeout(5000); - } + // Wait for all potential queue jobs + await allSettled(); + // And then close all queues await Promise.all([ this.systemQueue.close(), this.endedPollNotificationQueue.close(), diff --git a/packages/backend/src/core/ReactionService.ts b/packages/backend/src/core/ReactionService.ts index 3ca12551b1..2e8f76fa8a 100644 --- a/packages/backend/src/core/ReactionService.ts +++ b/packages/backend/src/core/ReactionService.ts @@ -28,6 +28,7 @@ import { UserBlockingService } from '@/core/UserBlockingService.js'; import { CustomEmojiService } from '@/core/CustomEmojiService.js'; import { RoleService } from '@/core/RoleService.js'; import { FeaturedService } from '@/core/FeaturedService.js'; +import { trackPromise } from '@/misc/promise-tracker.js'; const FALLBACK = '❤'; const PER_NOTE_REACTION_USER_PAIR_CACHE_MAX = 16; @@ -268,7 +269,7 @@ export class ReactionService { } } - dm.execute(); + trackPromise(dm.execute()); } //#endregion } @@ -316,7 +317,7 @@ export class ReactionService { dm.addDirectRecipe(reactee as MiRemoteUser); } dm.addFollowersRecipe(); - dm.execute(); + trackPromise(dm.execute()); } //#endregion } diff --git a/packages/backend/src/core/activitypub/ApDeliverManagerService.ts b/packages/backend/src/core/activitypub/ApDeliverManagerService.ts index 81003bcf1c..d7414e9c99 100644 --- a/packages/backend/src/core/activitypub/ApDeliverManagerService.ts +++ b/packages/backend/src/core/activitypub/ApDeliverManagerService.ts @@ -144,7 +144,7 @@ class DeliverManager { } // deliver - this.queueService.deliverMany(this.actor, this.activity, inboxes); + await this.queueService.deliverMany(this.actor, this.activity, inboxes); } } diff --git a/packages/backend/src/misc/promise-tracker.ts b/packages/backend/src/misc/promise-tracker.ts new file mode 100644 index 0000000000..c7166c6de9 --- /dev/null +++ b/packages/backend/src/misc/promise-tracker.ts @@ -0,0 +1,23 @@ +/* + * SPDX-FileCopyrightText: syuilo and other misskey contributors + * SPDX-License-Identifier: AGPL-3.0-only + */ + +const promiseRefs: Set>> = new Set(); + +/** + * This tracks promises that other modules decided not to wait for, + * and makes sure they are all settled before fully closing down the server. + */ +export function trackPromise(promise: Promise) { + if (process.env.NODE_ENV !== 'test') { + return; + } + const ref = new WeakRef(promise); + promiseRefs.add(ref); + promise.finally(() => promiseRefs.delete(ref)); +} + +export async function allSettled(): Promise { + await Promise.allSettled([...promiseRefs].map(r => r.deref())); +} diff --git a/packages/backend/src/server/api/endpoints/antennas/notes.ts b/packages/backend/src/server/api/endpoints/antennas/notes.ts index 0bf2688b4a..7293c2e39b 100644 --- a/packages/backend/src/server/api/endpoints/antennas/notes.ts +++ b/packages/backend/src/server/api/endpoints/antennas/notes.ts @@ -14,6 +14,7 @@ import { NoteEntityService } from '@/core/entities/NoteEntityService.js'; import { IdService } from '@/core/IdService.js'; import { FanoutTimelineService } from '@/core/FanoutTimelineService.js'; import { GlobalEventService } from '@/core/GlobalEventService.js'; +import { trackPromise } from '@/misc/promise-tracker.js'; import { ApiError } from '../../error.js'; export const meta = { @@ -92,7 +93,7 @@ export default class extends Endpoint { // eslint- antenna.isActive = true; antenna.lastUsedAt = new Date(); - this.antennasRepository.update(antenna.id, antenna); + trackPromise(this.antennasRepository.update(antenna.id, antenna)); if (needPublishEvent) { this.globalEventService.publishInternalEvent('antennaUpdated', antenna);