fix(server): DriveFile related N+1 query when call note packMany (again) (#10190)

* Revert "Revert "fix(server): DriveFile related N+1 query when call note packMany (#10133)""

This reverts commit a7c82eeabcc732e76c5c358c98812cac8457d57f.

* packManyByIdsMap: 存在チェックをしてなかったものは null を入れるように

* Note.packMany で reply とか renote がもうあったらそのファイルも引く

* テストを書く

* fix test

* fix test

* fix test

* fix test

3 files changed, 51 insertions, 4 deletions

diff --git a/packages/backend/src/core/entities/DriveFileEntityService.ts b/packages/backend/src/core/entities/DriveFileEntityService.ts
index 158fafa9d5..f769ddd5e9 100644
--- a/packages/backend/src/core/entities/DriveFileEntityService.ts
+++ b/packages/backend/src/core/entities/DriveFileEntityService.ts
@@ -1,5 +1,5 @@
 import { forwardRef, Inject, Injectable } from '@nestjs/common';
-import { DataSource } from 'typeorm';
+import { DataSource, In } from 'typeorm';
 import { DI } from '@/di-symbols.js';
 import type { NotesRepository, DriveFilesRepository } from '@/models/index.js';
 import type { Config } from '@/config.js';
@@ -21,6 +21,7 @@ type PackOptions = {
 };
 import { bindThis } from '@/decorators.js';
 import { isMimeImage } from '@/misc/is-mime-image.js';
+import { isNotNull } from '@/misc/is-not-null.js';
 
 @Injectable()
 export class DriveFileEntityService {
@@ -255,10 +256,33 @@ export class DriveFileEntityService {
 
 	@bindThis
 	public async packMany(
-		files: (DriveFile['id'] | DriveFile)[],
+		files: DriveFile[],
 		options?: PackOptions,
 	): Promise<Packed<'DriveFile'>[]> {
 		const items = await Promise.all(files.map(f => this.packNullable(f, options)));
 		return items.filter((x): x is Packed<'DriveFile'> => x != null);
 	}
+
+	@bindThis
+	public async packManyByIdsMap(
+		fileIds: DriveFile['id'][],
+		options?: PackOptions,
+	): Promise<Map<Packed<'DriveFile'>['id'], Packed<'DriveFile'> | null>> {
+		const files = await this.driveFilesRepository.findBy({ id: In(fileIds) });
+		const packedFiles = await this.packMany(files, options);
+		const map = new Map<Packed<'DriveFile'>['id'], Packed<'DriveFile'> | null>(packedFiles.map(f => [f.id, f]));
+		for (const id of fileIds) {
+			if (!map.has(id)) map.set(id, null);
+		}
+		return map;
+	}
+
+	@bindThis
+	public async packManyByIds(
+		fileIds: DriveFile['id'][],
+		options?: PackOptions,
+	): Promise<Packed<'DriveFile'>[]> {
+		const filesMap = await this.packManyByIdsMap(fileIds, options);
+		return fileIds.map(id => filesMap.get(id)).filter(isNotNull);
+	}
 }
diff --git a/packages/backend/src/core/entities/GalleryPostEntityService.ts b/packages/backend/src/core/entities/GalleryPostEntityService.ts
index ab29e7dba1..fb147ae181 100644
--- a/packages/backend/src/core/entities/GalleryPostEntityService.ts
+++ b/packages/backend/src/core/entities/GalleryPostEntityService.ts
@@ -41,7 +41,8 @@ export class GalleryPostEntityService {
 			title: post.title,
 			description: post.description,
 			fileIds: post.fileIds,
-			files: this.driveFileEntityService.packMany(post.fileIds),
+			// TODO: packMany causes N+1 queries
+			files: this.driveFileEntityService.packManyByIds(post.fileIds),
 			tags: post.tags.length > 0 ? post.tags : undefined,
 			isSensitive: post.isSensitive,
 			likedCount: post.likedCount,
diff --git a/packages/backend/src/core/entities/NoteEntityService.ts b/packages/backend/src/core/entities/NoteEntityService.ts
index 2ffe5f1c21..4ec10df9a6 100644
--- a/packages/backend/src/core/entities/NoteEntityService.ts
+++ b/packages/backend/src/core/entities/NoteEntityService.ts
@@ -11,6 +11,7 @@ import type { Note } from '@/models/entities/Note.js';
 import type { NoteReaction } from '@/models/entities/NoteReaction.js';
 import type { UsersRepository, NotesRepository, FollowingsRepository, PollsRepository, PollVotesRepository, NoteReactionsRepository, ChannelsRepository, DriveFilesRepository } from '@/models/index.js';
 import { bindThis } from '@/decorators.js';
+import { isNotNull } from '@/misc/is-not-null.js';
 import type { OnModuleInit } from '@nestjs/common';
 import type { CustomEmojiService } from '../CustomEmojiService.js';
 import type { ReactionService } from '../ReactionService.js';
@@ -249,6 +250,21 @@ export class NoteEntityService implements OnModuleInit {
 	}
 
 	@bindThis
+	public async packAttachedFiles(fileIds: Note['fileIds'], packedFiles: Map<Note['fileIds'][number], Packed<'DriveFile'> | null>): Promise<Packed<'DriveFile'>[]> {
+		const missingIds = [];
+		for (const id of fileIds) {
+			if (!packedFiles.has(id)) missingIds.push(id);
+		}
+		if (missingIds.length) {
+			const additionalMap = await this.driveFileEntityService.packManyByIdsMap(missingIds);
+			for (const [k, v] of additionalMap) {
+				packedFiles.set(k, v);
+			}
+		}
+		return fileIds.map(id => packedFiles.get(id)).filter(isNotNull);
+	}
+
+	@bindThis
 	public async pack(
 		src: Note['id'] | Note,
 		me?: { id: User['id'] } | null | undefined,
@@ -257,6 +273,7 @@ export class NoteEntityService implements OnModuleInit {
 			skipHide?: boolean;
 			_hint_?: {
 				myReactions: Map<Note['id'], NoteReaction | null>;
+				packedFiles: Map<Note['fileIds'][number], Packed<'DriveFile'> | null>;
 			};
 		},
 	): Promise<Packed<'Note'>> {
@@ -284,6 +301,7 @@ export class NoteEntityService implements OnModuleInit {
 		const reactionEmojiNames = Object.keys(note.reactions)
 			.filter(x => x.startsWith(':') && x.includes('@') && !x.includes('@.')) // リモートカスタム絵文字のみ
 			.map(x => this.reactionService.decodeReaction(x).reaction.replaceAll(':', ''));
+		const packedFiles = options?._hint_?.packedFiles;
 
 		const packed: Packed<'Note'> = await awaitAll({
 			id: note.id,
@@ -304,7 +322,7 @@ export class NoteEntityService implements OnModuleInit {
 			emojis: host != null ? this.customEmojiService.populateEmojis(note.emojis, host) : undefined,
 			tags: note.tags.length > 0 ? note.tags : undefined,
 			fileIds: note.fileIds,
-			files: this.driveFileEntityService.packMany(note.fileIds),
+			files: packedFiles != null ? this.packAttachedFiles(note.fileIds, packedFiles) : this.driveFileEntityService.packManyByIds(note.fileIds),
 			replyId: note.replyId,
 			renoteId: note.renoteId,
 			channelId: note.channelId ?? undefined,
@@ -388,11 +406,15 @@ export class NoteEntityService implements OnModuleInit {
 		}
 
 		await this.customEmojiService.prefetchEmojis(this.customEmojiService.aggregateNoteEmojis(notes));
+		// TODO: 本当は renote とか reply がないのに renoteId とか replyId があったらここで解決しておく
+		const fileIds = notes.map(n => [n.fileIds, n.renote?.fileIds, n.reply?.fileIds]).flat(2).filter(isNotNull);
+		const packedFiles = await this.driveFileEntityService.packManyByIdsMap(fileIds);
 
 		return await Promise.all(notes.map(n => this.pack(n, me, {
 			...options,
 			_hint_: {
 				myReactions: myReactionsMap,
+				packedFiles,
 			},
 		})));
 	}