diff --git a/packages/backend/src/core/entities/DriveFileEntityService.ts b/packages/backend/src/core/entities/DriveFileEntityService.ts index a6f7f369a6..1865d494c4 100644 --- a/packages/backend/src/core/entities/DriveFileEntityService.ts +++ b/packages/backend/src/core/entities/DriveFileEntityService.ts @@ -17,6 +17,7 @@ import { deepClone } from '@/misc/clone.js'; import { bindThis } from '@/decorators.js'; import { isMimeImage } from '@/misc/is-mime-image.js'; import { IdService } from '@/core/IdService.js'; +import { uniqueByKey } from '@/misc/unique-by-key.js'; import { UtilityService } from '../UtilityService.js'; import { VideoProcessingService } from '../VideoProcessingService.js'; import { UserEntityService } from './UserEntityService.js'; @@ -226,6 +227,7 @@ export class DriveFileEntityService { options?: PackOptions, hint?: { packedUser?: Packed<'UserLite'> + packedFolder?: Packed<'DriveFolder'> }, ): Promise | null> { const opts = Object.assign({ @@ -250,9 +252,9 @@ export class DriveFileEntityService { thumbnailUrl: this.getThumbnailUrl(file), comment: file.comment, folderId: file.folderId, - folder: opts.detail && file.folderId ? this.driveFolderEntityService.pack(file.folderId, { + folder: opts.detail && file.folderId ? (hint?.packedFolder ?? this.driveFolderEntityService.pack(file.folderId, { detail: true, - }) : null, + })) : null, userId: file.userId, user: (opts.withUser && file.userId) ? hint?.packedUser ?? this.userEntityService.pack(file.userId) : null, }); @@ -263,10 +265,41 @@ export class DriveFileEntityService { files: MiDriveFile[], options?: PackOptions, ): Promise[]> { - const _user = files.map(({ user, userId }) => user ?? userId).filter(x => x != null); - const _userMap = await this.userEntityService.packMany(_user) - .then(users => new Map(users.map(user => [user.id, user]))); - const items = await Promise.all(files.map(f => this.packNullable(f, options, f.userId ? { packedUser: _userMap.get(f.userId) } : {}))); + // -- ユーザ情報の事前取得 -- + + let userMap: Map> | null = null; + if (options?.withUser) { + const users = files + .map(({ user, userId }) => user ?? userId) + .filter(x => x != null); + + const uniqueUsers = uniqueByKey(users, (user) => typeof user === 'string' ? user : user.id); + const packedUsers = await this.userEntityService.packMany(uniqueUsers); + userMap = new Map(packedUsers.map(user => [user.id, user])); + } + + // -- フォルダ情報の事前取得 -- + + let folderMap: Map> | null = null; + if (options?.detail) { + const folders = files + .map(({ folder, folderId }) => folder ?? folderId) + .filter(x => x != null); + + const uniqueFolders = uniqueByKey(folders, (folder) => typeof folder === 'string' ? folder : folder.id); + const packedFolders = await this.driveFolderEntityService.packMany(uniqueFolders, { detail: true }); + folderMap = new Map(packedFolders.map(folder => [folder.id, folder])); + } + + const items = await Promise.all(files.map(f => this.packNullable( + f, + options, + { + packedUser: f.userId ? userMap?.get(f.userId) : undefined, + packedFolder: f.folderId ? folderMap?.get(f.folderId) : undefined, + }, + ))); + return items.filter(x => x != null); } diff --git a/packages/backend/src/core/entities/DriveFolderEntityService.ts b/packages/backend/src/core/entities/DriveFolderEntityService.ts index 299f23ad38..326421e149 100644 --- a/packages/backend/src/core/entities/DriveFolderEntityService.ts +++ b/packages/backend/src/core/entities/DriveFolderEntityService.ts @@ -12,6 +12,9 @@ import type { } from '@/models/Blocking.js'; import type { MiDriveFolder } from '@/models/DriveFolder.js'; import { bindThis } from '@/decorators.js'; import { IdService } from '@/core/IdService.js'; +import { In } from 'typeorm'; +import { uniqueByKey } from '@/misc/unique-by-key.js'; +import { splitIdAndObjects } from '@/misc/split-id-and-objects.js'; @Injectable() export class DriveFolderEntityService { @@ -32,12 +35,20 @@ export class DriveFolderEntityService { options?: { detail: boolean }, + hint?: { + folderMap?: Map; + foldersCountMap?: Map | null; + filesCountMap?: Map | null; + parentPacker?: (id: string) => Promise>; + }, ): Promise> { const opts = Object.assign({ detail: false, }, options); - const folder = typeof src === 'object' ? src : await this.driveFoldersRepository.findOneByOrFail({ id: src }); + const folder = typeof src === 'object' + ? src + : hint?.folderMap?.get(src) ?? await this.driveFoldersRepository.findOneByOrFail({ id: src }); return await awaitAll({ id: folder.id, @@ -46,20 +57,141 @@ export class DriveFolderEntityService { parentId: folder.parentId, ...(opts.detail ? { - foldersCount: this.driveFoldersRepository.countBy({ - parentId: folder.id, - }), - filesCount: this.driveFilesRepository.countBy({ - folderId: folder.id, - }), + foldersCount: hint?.foldersCountMap?.get(folder.id) + ?? this.driveFoldersRepository.countBy({ + parentId: folder.id, + }), + filesCount: hint?.filesCountMap?.get(folder.id) + ?? this.driveFilesRepository.countBy({ + folderId: folder.id, + }), ...(folder.parentId ? { - parent: this.pack(folder.parentId, { - detail: true, - }), + parent: hint?.parentPacker + ? hint.parentPacker(folder.parentId) + : this.pack(folder.parentId, { detail: true }, hint), } : {}), } : {}), }); } -} + public async packMany( + src: Array, + options?: { + detail: boolean + }, + ): Promise>> { + /** + * 重複を除去しつつ、必要なDriveFolderオブジェクトをすべて取得する + */ + const collectUniqueObjects = async (src: Array) => { + const uniqueSrc = uniqueByKey( + src, + (s) => typeof s === 'string' ? s : s.id, + ); + const { ids, objects } = splitIdAndObjects(uniqueSrc); + + const uniqueObjects = new Map(objects.map(s => [s.id, s])); + const needsFetchIds = ids.filter(id => !uniqueObjects.has(id)); + + if (needsFetchIds.length > 0) { + const fetchedObjects = await this.driveFoldersRepository.find({ + where: { + id: In(needsFetchIds), + }, + }); + for (const obj of fetchedObjects) { + uniqueObjects.set(obj.id, obj); + } + } + + return uniqueObjects; + }; + + /** + * 親フォルダーを再帰的に収集する + */ + const collectAncestors = async (folderMap: Map) => { + for (;;) { + const parentIds = new Set(); + for (const folder of folderMap.values()) { + if (folder.parentId != null && !folderMap.has(folder.parentId)) { + parentIds.add(folder.parentId); + } + } + + if (parentIds.size === 0) break; + + const fetchedParents = await this.driveFoldersRepository.find({ + where: { + id: In([...parentIds]), + }, + }); + + if (fetchedParents.length === 0) break; + + for (const parent of fetchedParents) { + folderMap.set(parent.id, parent); + } + } + }; + + const opts = Object.assign({ + detail: false, + }, options); + + const folderMap = await collectUniqueObjects(src); + + let foldersCountMap: Map | null = null; + let filesCountMap: Map | null = null; + if (opts.detail) { + await collectAncestors(folderMap); + + const ids = [...folderMap.keys()]; + if (ids.length > 0) { + const folderCounts = await this.driveFoldersRepository.createQueryBuilder('folder') + .select('folder.parentId', 'parentId') + .addSelect('COUNT(*)', 'count') + .where('folder.parentId IN (:...ids)', { ids }) + .groupBy('folder.parentId') + .getRawMany<{ parentId: string; count: string }>(); + + const fileCounts = await this.driveFilesRepository.createQueryBuilder('file') + .select('file.folderId', 'folderId') + .addSelect('COUNT(*)', 'count') + .where('file.folderId IN (:...ids)', { ids }) + .groupBy('file.folderId') + .getRawMany<{ folderId: string; count: string }>(); + + foldersCountMap = new Map(folderCounts.map(row => [row.parentId, Number(row.count)])); + filesCountMap = new Map(fileCounts.map(row => [row.folderId, Number(row.count)])); + } else { + foldersCountMap = new Map(); + filesCountMap = new Map(); + } + } + + const packedMap = new Map>>(); + const packFromId = (id: string): Promise> => { + const cached = packedMap.get(id); + if (cached) return cached; + + const folder = folderMap.get(id); + if (!folder) { + throw new Error(`DriveFolder not found: ${id}`); + } + + const packedPromise = this.pack(folder, options, { + folderMap, + foldersCountMap, + filesCountMap, + parentPacker: packFromId, + }); + packedMap.set(id, packedPromise); + + return packedPromise; + }; + + return Promise.all(src.map(s => packFromId(typeof s === 'string' ? s : s.id))); + } +} diff --git a/packages/backend/src/misc/split-id-and-objects.ts b/packages/backend/src/misc/split-id-and-objects.ts new file mode 100644 index 0000000000..d23bb93695 --- /dev/null +++ b/packages/backend/src/misc/split-id-and-objects.ts @@ -0,0 +1,27 @@ +/* + * SPDX-FileCopyrightText: syuilo and misskey-project + * SPDX-License-Identifier: AGPL-3.0-only + */ + +/** + * idとオブジェクトを分離する + * @param input idまたはオブジェクトの配列 + * @returns idの配列とオブジェクトの配列 + */ +export function splitIdAndObjects(input: (T | string)[]): { ids: string[]; objects: T[] } { + const ids: string[] = []; + const objects : T[] = []; + + for (const item of input) { + if (typeof item === 'string') { + ids.push(item); + } else { + objects.push(item); + } + } + + return { + ids, + objects, + }; +} diff --git a/packages/backend/src/misc/unique-by-key.ts b/packages/backend/src/misc/unique-by-key.ts new file mode 100644 index 0000000000..4308e29d21 --- /dev/null +++ b/packages/backend/src/misc/unique-by-key.ts @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: syuilo and misskey-project + * SPDX-License-Identifier: AGPL-3.0-only + */ + +/** + * itemsの中でkey関数が返す値が重複しないようにした配列を返す + * @param items 重複を除去したい配列 + * @param key 重複判定に使うキーを返す関数 + * @returns 重複を除去した配列 + */ +export function uniqueByKey(items: Iterable, key: (item: TItem) => TKey): TItem[] { + const map = new Map(); + for (const item of items) { + const k = key(item); + if (!map.has(k)) { + map.set(k, item); + } + } + return [...map.values()]; +} diff --git a/packages/backend/test/unit/entities/DriveFileEntityService.ts b/packages/backend/test/unit/entities/DriveFileEntityService.ts new file mode 100644 index 0000000000..2e416326ee --- /dev/null +++ b/packages/backend/test/unit/entities/DriveFileEntityService.ts @@ -0,0 +1,227 @@ +/* + * SPDX-FileCopyrightText: syuilo and misskey-project + * SPDX-License-Identifier: AGPL-3.0-only + */ + +process.env.NODE_ENV = 'test'; + +import { afterAll, beforeAll, beforeEach, describe, expect, jest, test } from '@jest/globals'; +import { Test } from '@nestjs/testing'; +import type { TestingModule } from '@nestjs/testing'; +import type { DriveFilesRepository, DriveFoldersRepository, UsersRepository } from '@/models/_.js'; +import { GlobalModule } from '@/GlobalModule.js'; +import { CoreModule } from '@/core/CoreModule.js'; +import { DriveFileEntityService } from '@/core/entities/DriveFileEntityService.js'; +import { DriveFolderEntityService } from '@/core/entities/DriveFolderEntityService.js'; +import { UserEntityService } from '@/core/entities/UserEntityService.js'; +import { DI } from '@/di-symbols.js'; +import { genAidx } from '@/misc/id/aidx.js'; +import { secureRndstr } from '@/misc/secure-rndstr.js'; + +const describeBenchmark = process.env.RUN_BENCHMARKS === '1' ? describe : describe.skip; + +describe('DriveFileEntityService', () => { + let app: TestingModule; + let service: DriveFileEntityService; + let driveFolderEntityService: DriveFolderEntityService; + let driveFilesRepository: DriveFilesRepository; + let driveFoldersRepository: DriveFoldersRepository; + let usersRepository: UsersRepository; + let idCounter = 0; + + const userEntityServiceMock = { + packMany: jest.fn(async (users: Array) => { + return users.map(u => ({ + id: typeof u === 'string' ? u : u.id, + username: 'user', + })); + }), + pack: jest.fn(async (user: string | { id: string }) => { + return { + id: typeof user === 'string' ? user : user.id, + username: 'user', + }; + }), + }; + + const nextId = () => genAidx(Date.now() + (idCounter++)); + + const createUser = async () => { + const un = secureRndstr(16); + const id = nextId(); + await usersRepository.insert({ + id, + username: un, + usernameLower: un.toLowerCase(), + }); + return usersRepository.findOneByOrFail({ id }); + }; + + const createFolder = async (name: string, parentId: string | null) => { + const id = nextId(); + await driveFoldersRepository.insert({ + id, + name, + userId: null, + parentId, + }); + return driveFoldersRepository.findOneByOrFail({ id }); + }; + + const createFile = async (folderId: string | null, userId: string | null) => { + const id = nextId(); + await driveFilesRepository.insert({ + id, + userId, + userHost: null, + md5: secureRndstr(32), + name: `file-${id}`, + type: 'text/plain', + size: 1, + comment: null, + blurhash: null, + properties: {}, + storedInternal: true, + url: `https://example.com/${id}`, + thumbnailUrl: null, + webpublicUrl: null, + webpublicType: null, + accessKey: null, + thumbnailAccessKey: null, + webpublicAccessKey: null, + uri: null, + src: null, + folderId, + isSensitive: false, + maybeSensitive: false, + maybePorn: false, + isLink: false, + requestHeaders: null, + requestIp: null, + }); + return driveFilesRepository.findOneByOrFail({ id }); + }; + + beforeAll(async () => { + const moduleBuilder = Test.createTestingModule({ + imports: [GlobalModule, CoreModule], + }); + moduleBuilder.overrideProvider(UserEntityService).useValue(userEntityServiceMock as any); + + app = await moduleBuilder.compile(); + await app.init(); + app.enableShutdownHooks(); + + service = app.get(DriveFileEntityService); + driveFolderEntityService = app.get(DriveFolderEntityService); + driveFilesRepository = app.get(DI.driveFilesRepository); + driveFoldersRepository = app.get(DI.driveFoldersRepository); + usersRepository = app.get(DI.usersRepository); + }); + + beforeEach(() => { + userEntityServiceMock.packMany.mockClear(); + userEntityServiceMock.pack.mockClear(); + }); + + afterAll(async () => { + await app.close(); + }); + + describe('pack', () => { + test('detail: false', async () => { + const user = await createUser(); + const folder = await createFolder('pack-root', null); + const file = await createFile(folder.id, user.id); + + const packed = await service.pack(file, { detail: false, self: true }) as any; + expect(packed.id).toBe(file.id); + expect(packed.folder).toBeNull(); + expect(packed.user).toBeNull(); + expect(packed.userId).toBeNull(); + }); + + test('detail: true', async () => { + const folder = await createFolder('pack-parent', null); + const child = await createFolder('pack-child', folder.id); + const file = await createFile(child.id, null); + + const packed = await service.pack(file, { detail: true, self: true }) as any; + expect(packed.folder?.id).toBe(child.id); + expect(packed.folder?.parent?.id).toBe(folder.id); + }); + }); + + describe('packNullable', () => { + test('returns null for missing', async () => { + const packed = await service.packNullable('non-existent' as any, { detail: false }); + expect(packed).toBeNull(); + }); + + test('uses packedUser hint when withUser', async () => { + const user = await createUser(); + const file = await createFile(null, user.id); + + const packed = await service.packNullable(file, { withUser: true, self: true }, { + packedUser: { id: user.id, username: 'hint' } as any, + }); + expect(packed?.user?.id).toBe(user.id); + expect(packed?.user?.username).toBe('hint'); + }); + }); + + describe('packMany', () => { + test('withUser: true uses deduped packMany', async () => { + const user = await createUser(); + const fileA = await createFile(null, user.id); + const fileB = await createFile(null, user.id); + + const packed = await service.packMany([fileA, fileB], { withUser: true, self: true }); + expect(packed.length).toBe(2); + expect(userEntityServiceMock.packMany).toHaveBeenCalledTimes(1); + expect(userEntityServiceMock.packMany.mock.calls[0]?.[0]?.length).toBe(1); + expect(packed[0]?.user?.id).toBe(user.id); + }); + + test('detail: true packs folder', async () => { + const folder = await createFolder('packmany-root', null); + const file = await createFile(folder.id, null); + + const packed = await service.packMany([file], { detail: true, self: true }); + expect(packed[0]?.folder?.id).toBe(folder.id); + expect(packed[0]?.folder?.parent).toBeUndefined(); + }); + + test('detail: true uses DriveFolderEntityService pack', async () => { + const folder = await createFolder('packmany-folder', null); + const file = await createFile(folder.id, null); + const packSpy = jest.spyOn(driveFolderEntityService, 'pack'); + + await service.packMany([file], { detail: true, self: true }); + expect(packSpy).toHaveBeenCalled(); + packSpy.mockRestore(); + }); + }); + + describeBenchmark('benchmark', () => { + test('packMany', async () => { + const user = await createUser(); + const folders = []; + for (let i = 0; i < 100; i++) { + folders.push(await createFolder(`bench-${i}`, null)); + } + const files = []; + for (const folder of folders) { + for (let j = 0; j < 20; j++) { + files.push(await createFile(folder.id, user.id)); + } + } + + const start = Date.now(); + await service.packMany(files, { detail: true, withUser: true, self: true }); + const elapsed = Date.now() - start; + + console.log(`DriveFileEntityService.packMany benchmark: ${elapsed}ms`); + }); + }); +}); diff --git a/packages/backend/test/unit/entities/DriveFolderEntityService.ts b/packages/backend/test/unit/entities/DriveFolderEntityService.ts new file mode 100644 index 0000000000..299ee5f42b --- /dev/null +++ b/packages/backend/test/unit/entities/DriveFolderEntityService.ts @@ -0,0 +1,171 @@ +/* + * SPDX-FileCopyrightText: syuilo and misskey-project + * SPDX-License-Identifier: AGPL-3.0-only + */ + +process.env.NODE_ENV = 'test'; + +import { afterAll, beforeAll, describe, expect, test } from '@jest/globals'; +import { Test } from '@nestjs/testing'; +import type { TestingModule } from '@nestjs/testing'; +import type { DriveFilesRepository, DriveFoldersRepository } from '@/models/_.js'; +import { GlobalModule } from '@/GlobalModule.js'; +import { CoreModule } from '@/core/CoreModule.js'; +import { DriveFolderEntityService } from '@/core/entities/DriveFolderEntityService.js'; +import { DI } from '@/di-symbols.js'; +import { genAidx } from '@/misc/id/aidx.js'; +import { secureRndstr } from '@/misc/secure-rndstr.js'; + +const describeBenchmark = process.env.RUN_BENCHMARKS === '1' ? describe : describe.skip; + +describe('DriveFolderEntityService', () => { + let app: TestingModule; + let service: DriveFolderEntityService; + let driveFoldersRepository: DriveFoldersRepository; + let driveFilesRepository: DriveFilesRepository; + let idCounter = 0; + + const nextId = () => genAidx(Date.now() + (idCounter++)); + + const createFolder = async (name: string, parentId: string | null) => { + const id = nextId(); + await driveFoldersRepository.insert({ + id, + name, + userId: null, + parentId, + }); + return driveFoldersRepository.findOneByOrFail({ id }); + }; + + const createFile = async (folderId: string | null) => { + const id = nextId(); + await driveFilesRepository.insert({ + id, + userId: null, + userHost: null, + md5: secureRndstr(32), + name: `file-${id}`, + type: 'text/plain', + size: 1, + comment: null, + blurhash: null, + properties: {}, + storedInternal: true, + url: `https://example.com/${id}`, + thumbnailUrl: null, + webpublicUrl: null, + webpublicType: null, + accessKey: null, + thumbnailAccessKey: null, + webpublicAccessKey: null, + uri: null, + src: null, + folderId, + isSensitive: false, + maybeSensitive: false, + maybePorn: false, + isLink: false, + requestHeaders: null, + requestIp: null, + }); + }; + + beforeAll(async () => { + app = await Test.createTestingModule({ + imports: [GlobalModule, CoreModule], + }).compile(); + await app.init(); + app.enableShutdownHooks(); + + service = app.get(DriveFolderEntityService); + driveFoldersRepository = app.get(DI.driveFoldersRepository); + driveFilesRepository = app.get(DI.driveFilesRepository); + }); + + afterAll(async () => { + await app.close(); + }); + + describe('pack', () => { + test('detail: false', async () => { + const root = await createFolder('root', null); + const child = await createFolder('child', root.id); + + const packed = await service.pack(child, { detail: false }) as any; + expect(packed.id).toBe(child.id); + expect(packed.parentId).toBe(root.id); + expect(packed.parent).toBeUndefined(); + expect(packed.foldersCount).toBeUndefined(); + expect(packed.filesCount).toBeUndefined(); + }); + + test('detail: true', async () => { + const root = await createFolder('root-detail', null); + const child = await createFolder('child-detail', root.id); + await createFolder('grandchild-detail', child.id); + await createFile(child.id); + await createFile(child.id); + + const packed = await service.pack(child, { detail: true }) as any; + expect(packed.id).toBe(child.id); + expect(packed.foldersCount).toBe(1); + expect(packed.filesCount).toBe(2); + expect(packed.parent?.id).toBe(root.id); + expect(packed.parent?.parent).toBeUndefined(); + }); + + test('detail: true reaches root for deep hierarchy', async () => { + const root = await createFolder('root-deep', null); + const level1 = await createFolder('level-1', root.id); + const level2 = await createFolder('level-2', level1.id); + const level3 = await createFolder('level-3', level2.id); + const level4 = await createFolder('level-4', level3.id); + const level5 = await createFolder('level-5', level4.id); + + const packed = await service.pack(level5, { detail: true }) as any; + expect(packed.id).toBe(level5.id); + expect(packed.parent?.id).toBe(level4.id); + expect(packed.parent?.parent?.id).toBe(level3.id); + expect(packed.parent?.parent?.parent?.id).toBe(level2.id); + expect(packed.parent?.parent?.parent?.parent?.id).toBe(level1.id); + expect(packed.parent?.parent?.parent?.parent?.parent?.id).toBe(root.id); + expect(packed.parent?.parent?.parent?.parent?.parent?.parent).toBeUndefined(); + }); + }); + + describe('packMany', () => { + test('preserves order and packs parents', async () => { + const root = await createFolder('root-many', null); + const childA = await createFolder('child-a', root.id); + const childB = await createFolder('child-b', root.id); + await createFolder('child-a-sub', childA.id); + await createFile(childA.id); + + const packed = await service.packMany([childB, childA], { detail: true }) as any; + expect(packed[0].id).toBe(childB.id); + expect(packed[1].id).toBe(childA.id); + expect(packed[0].parent?.id).toBe(root.id); + expect(packed[1].parent?.id).toBe(root.id); + expect(packed[0].filesCount).toBe(0); + expect(packed[1].filesCount).toBe(1); + expect(packed[0].foldersCount).toBe(0); + expect(packed[1].foldersCount).toBe(1); + }); + }); + + describeBenchmark('benchmark', () => { + test('packMany', async () => { + const root = await createFolder('bench-root', null); + const folders = []; + for (let i = 0; i < 200; i++) { + folders.push(await createFolder(`bench-${i}`, root.id)); + } + + const start = Date.now(); + await service.packMany(folders, { detail: true }); + const elapsed = Date.now() - start; + console.log(`DriveFolderEntityService.packMany benchmark: ${elapsed}ms`); + }); + }); +});