spec(backend/drive/files): ネットワーク不安定・高負荷時ファイルが重複してアップロードされる問題を修正 (MisskeyIO#473)

Co-authored-by: riku6460 <17585784+riku6460@users.noreply.github.com>

packages/backend/src/server/api/endpoints/drive/files/create.ts

@@ -3,17 +3,22 @@
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
+import * as fs from 'node:fs';
+import { createHash } from 'crypto';
+import * as stream from 'node:stream/promises';
+import * as Redis from 'ioredis';
 import ms from 'ms';
-import { Injectable } from '@nestjs/common';
-import { DB_MAX_IMAGE_COMMENT_LENGTH } from '@/const.js';
-import type Logger from '@/logger.js';
-import { IdentifiableError } from '@/misc/identifiable-error.js';
-import { Endpoint } from '@/server/api/endpoint-base.js';
+import { Inject, Injectable } from '@nestjs/common';
+import type { DriveFilesRepository } from '@/models/_.js';
+import { DI } from '@/di-symbols.js';
 import { DriveFileEntityService } from '@/core/entities/DriveFileEntityService.js';
-import { MetaService } from '@/core/MetaService.js';
-import { LoggerService } from '@/core/LoggerService.js';
 import { DriveService } from '@/core/DriveService.js';
-import { ApiError } from '../../../error.js';
+import { Endpoint } from '@/server/api/endpoint-base.js';
+import { IdentifiableError } from '@/misc/identifiable-error.js';
+import { DB_MAX_IMAGE_COMMENT_LENGTH } from '@/const.js';
+import { LoggerService } from '@/core/LoggerService.js';
+import { MetaService } from '@/core/MetaService.js';
+import { ApiError } from '@/server/api/error.js';
 
 export const meta = {
 	tags: ['drive'],
@@ -31,7 +36,6 @@ export const meta = {
 	requireFile: true,
 
 	kind: 'write:drive',
-
 	description: 'Upload a new drive file.',
 
 	res: {
@@ -41,6 +45,19 @@ export const meta = {
 	},
 
 	errors: {
+		invalidParam: {
+			message: 'Invalid param.',
+			code: 'INVALID_PARAM',
+			id: 'b2da4a73-a9d2-44e5-a81b-28e796874fc3',
+		},
+
+		processing: {
+			message: 'We are processing your request. Please wait a moment.',
+			code: 'PROCESSING',
+			id: 'b495d816-b077-4dc1-b135-7fde73fcca5e',
+			httpStatusCode: 202,
+		},
+
 		invalidFileName: {
 			message: 'Invalid file name.',
 			code: 'INVALID_FILE_NAME',
@@ -75,17 +92,54 @@ export const paramDef = {
 
 @Injectable()
 export default class extends Endpoint<typeof meta, typeof paramDef> { // eslint-disable-line import/no-default-export
-	private logger: Logger;
-
 	constructor(
+		@Inject(DI.redis)
+		private redisClient: Redis.Redis,
+
+		@Inject(DI.driveFilesRepository)
+		private driveFilesRepository: DriveFilesRepository,
+
 		private driveFileEntityService: DriveFileEntityService,
 		private metaService: MetaService,
 		private loggerService: LoggerService,
 		private driveService: DriveService,
 	) {
 		super(meta, paramDef, async (ps, me, _, file, cleanup, ip, headers) => {
+			const logger = this.loggerService.getLogger('api:drive:files:create');
+
+			if (!file) {
+				logger.setContext({ userId: me.id, ip, headers });
+				logger.error('File is required but did not provided.');
+				throw new ApiError(meta.errors.invalidParam);
+			}
+
+			const calcHash = createHash('sha256').update(`${ps.folderId}:${ps.isSensitive}`);
+			await stream.pipeline(fs.createReadStream(file.path, { encoding: 'binary', start: 0, end: 1024 * 1024 }), calcHash);
+			const hash = calcHash.digest('base64');
+			logger.setContext({ userId: me.id, hash, ip, headers });
+			logger.info('Request to create drive file.');
+
+			const idempotent = process.env.FORCE_IGNORE_IDEMPOTENCY_FOR_TESTING !== 'true' ? await this.redisClient.get(`drive:files:create:idempotent:${me.id}:${hash}`) : null;
+			if (idempotent === '_') { // 他のサーバーで処理中
+				logger.warn('The request is being processed by another server.');
+				throw new ApiError(meta.errors.processing);
+			}
+
+			// すでに同じリクエストが処理されている場合、そのファイルを返す
+			// ただし、記録されているファイルが見つからない場合は、新規として処理を続行
+			if (idempotent) {
+				const driveFile = await this.driveFilesRepository.findOneBy({ id: idempotent });
+				if (driveFile) {
+					logger.info('The request has already been processed.', { fileId: driveFile.id });
+					return await this.driveFileEntityService.pack(driveFile, me, { self: true });
+				}
+			}
+
+			// 30秒の間、リクエストを処理中として記録
+			await this.redisClient.set(`drive:files:create:idempotent:${me.id}:${hash}`, '_', 'EX', 30);
+
 			// Get 'name' parameter
-			let name = ps.name ?? file!.name ?? null;
+			let name = ps.name ?? file.name ?? null;
 			if (name != null) {
 				name = name.trim();
 				if (name.length === 0) {
@@ -103,7 +157,7 @@ export default class extends Endpoint<typeof meta, typeof paramDef> { // eslint-
 				// Create file
 				const driveFile = await this.driveService.addFile({
 					user: me,
-					path: file!.path,
+					path: file.path,
 					name,
 					comment: ps.comment,
 					folderId: ps.folderId,
@@ -112,13 +166,23 @@ export default class extends Endpoint<typeof meta, typeof paramDef> { // eslint-
 					requestIp: instance.enableIpLogging ? ip : null,
 					requestHeaders: instance.enableIpLogging ? headers : null,
 				});
+
+				// 1分間、リクエストの処理結果を記録
+				await this.redisClient.set(`drive:files:create:idempotent:${me.id}:${hash}`, driveFile.id, 'EX', 60);
+
+				logger.info('Successfully created drive file.', { fileId: driveFile.id });
 				return await this.driveFileEntityService.pack(driveFile, me, { self: true });
 			} catch (e) {
-				this.logger.error('Failed to create drive file', { error: e });
+				// エラーが発生した場合、リクエストの処理結果を削除
+				await this.redisClient.unlink(`drive:files:create:idempotent:${me.id}:${hash}`);
+
+				logger.error('Failed to create drive file.', { error: e });
+
 				if (e instanceof IdentifiableError) {
 					if (e.id === '282f77bf-5816-4f72-9264-aa14d8261a21') throw new ApiError(meta.errors.inappropriate);
 					if (e.id === 'c6244ed2-a39a-4e1c-bf93-f0fbd7764fa6') throw new ApiError(meta.errors.noFreeSpace);
 				}
+
 				const err = e as Error;
 				throw new ApiError(
 					{
@@ -130,14 +194,12 @@ export default class extends Endpoint<typeof meta, typeof paramDef> { // eslint-
 						e: {
 							message: err.message,
 							code: err.name,
-						}
-					}
+						},
+					},
 				);
 			} finally {
-				cleanup!();
+				if (cleanup) cleanup();
 			}
 		});
-
-		this.logger = this.loggerService.getLogger('api:drive:files:create');
 	}
 }

packages/backend/src/server/api/endpoints/drive/files/upload-from-url.ts

@@ -3,29 +3,43 @@
  * SPDX-License-Identifier: AGPL-3.0-only
  */
 
+import { createHash } from 'crypto';
 import ms from 'ms';
-import { Injectable } from '@nestjs/common';
-import { Endpoint } from '@/server/api/endpoint-base.js';
-import { GlobalEventService } from '@/core/GlobalEventService.js';
+import * as Redis from 'ioredis';
+import { Inject, Injectable } from '@nestjs/common';
+import type { DriveFilesRepository } from '@/models/_.js';
+import { ApiError } from '@/server/api/error.js';
+import { DI } from '@/di-symbols.js';
 import { DriveFileEntityService } from '@/core/entities/DriveFileEntityService.js';
 import { DriveService } from '@/core/DriveService.js';
+import { Endpoint } from '@/server/api/endpoint-base.js';
+import { GlobalEventService } from '@/core/GlobalEventService.js';
+import { LoggerService } from '@/core/LoggerService.js';
 
 export const meta = {
 	tags: ['drive'],
 
-	limit: {
-		duration: ms('1hour'),
-		max: 60,
-	},
-
-	description: 'Request the server to download a new drive file from the specified URL.',
-
 	requireCredential: true,
 	requireRolePolicy: 'canCreateContent',
 
 	prohibitMoved: true,
 
+	limit: {
+		duration: ms('1hour'),
+		max: 60,
+	},
+
 	kind: 'write:drive',
+	description: 'Request the server to download a new drive file from the specified URL.',
+
+	errors: {
+		processing: {
+			message: 'We are processing your request. Please wait a moment.',
+			code: 'PROCESSING',
+			id: '59953963-7b7a-4b6a-b74b-2e9f86992b04',
+			httpStatusCode: 202,
+		},
+	},
 } as const;
 
 export const paramDef = {
@@ -44,19 +58,70 @@ export const paramDef = {
 @Injectable()
 export default class extends Endpoint<typeof meta, typeof paramDef> { // eslint-disable-line import/no-default-export
 	constructor(
-		private driveFileEntityService: DriveFileEntityService,
+		@Inject(DI.redis)
+		private redisClient: Redis.Redis,
+
+		@Inject(DI.driveFilesRepository)
+		private driveFilesRepository: DriveFilesRepository,
+
+		private loggerService: LoggerService,
 		private driveService: DriveService,
+		private driveFileEntityService: DriveFileEntityService,
 		private globalEventService: GlobalEventService,
 	) {
-		super(meta, paramDef, async (ps, user, _1, _2, _3, ip, headers) => {
-			this.driveService.uploadFromUrl({ url: ps.url, user, folderId: ps.folderId, sensitive: ps.isSensitive, force: ps.force, comment: ps.comment, requestIp: ip, requestHeaders: headers }).then(file => {
-				this.driveFileEntityService.pack(file, user, { self: true }).then(packedFile => {
-					this.globalEventService.publishMainStream(user.id, 'urlUploadFinished', {
-						marker: ps.marker,
-						file: packedFile,
+		super(meta, paramDef, async (ps, me, _token, _file, _cleanup, ip, headers) => {
+			const logger = this.loggerService.getLogger('api:drive:files:upload-from-url');
+			const hash = createHash('sha256').update(`${ps.folderId}:${ps.url}:${ps.isSensitive}`).digest('base64');
+			logger.setContext({ userId: me.id, hash, ip, headers });
+			logger.info('Request to upload from URL.');
+
+			const idempotent = process.env.FORCE_IGNORE_IDEMPOTENCY_FOR_TESTING !== 'true' ? await this.redisClient.get(`drive:files:upload-from-url:idempotent:${me.id}:${hash}`) : null;
+			if (idempotent === '_') { // 他のサーバーで処理中
+				logger.warn('The request is being processed by another server.');
+				throw new ApiError(meta.errors.processing);
+			}
+
+			// すでに同じリクエストが処理されている場合、そのファイルを返す
+			// ただし、記録されているファイルが見つからない場合は、新規として処理を続行
+			if (idempotent) {
+				const file = await this.driveFilesRepository.findOneBy({ id: idempotent });
+				if (file) {
+					logger.info('The request has already been processed.', { fileId: file.id });
+					return;
+				}
+			}
+
+			// 30秒の間、リクエストを処理中として記録
+			await this.redisClient.set(`drive:files:upload-from-url:idempotent:${me.id}:${hash}`, '_', 'EX', 30);
+
+			this.driveService.uploadFromUrl({
+				url: ps.url,
+				user: me,
+				folderId: ps.folderId,
+				sensitive: ps.isSensitive,
+				force: ps.force,
+				comment: ps.comment,
+				requestIp: ip,
+				requestHeaders: headers,
+			}).then(
+				async file => {
+					// 1分間、リクエストの処理結果を記録
+					await this.redisClient.set(`drive:files:upload-from-url:idempotent:${me.id}:${hash}`, file.id, 'EX', 60);
+					logger.info('Successfully uploaded from URL.', { fileId: file.id });
+
+					this.driveFileEntityService.pack(file, me, { self: true }).then(packedFile => {
+						this.globalEventService.publishMainStream(me.id, 'urlUploadFinished', {
+							marker: ps.marker,
+							file: packedFile,
+						});
 					});
-				});
-			});
+				},
+				async err => {
+					// エラーが発生した場合、リクエストの処理結果を削除
+					await this.redisClient.unlink(`drive:files:upload-from-url:idempotent:${me.id}:${hash}`);
+					logger.error('Failed to upload from URL.', { error: err });
+				},
+			);
 		});
 	}
 }

packages/backend/test/e2e/drive.ts

@@ -8,7 +8,7 @@ process.env.NODE_ENV = 'test';
 import * as assert from 'assert';
 import { MiNote } from '@/models/Note.js';
 import type { Packed } from '@/misc/json-schema.js';
-import { api, initTestDb, makeStreamCatcher, post, signup, uploadFile } from '../utils.js';
+import { api, initTestDb, makeStreamCatcher, post, sendEnvUpdateRequest, signup, uploadFile } from '../utils.js';
 import type * as misskey from 'misskey-js';
 import type{ Repository } from 'typeorm';
 
@@ -19,6 +19,8 @@ describe('Drive', () => {
 	let bob: misskey.entities.SignupResponse;
 
 	beforeAll(async () => {
+		await sendEnvUpdateRequest({ key: 'FORCE_IGNORE_IDEMPOTENCY_FOR_TESTING', value: 'true' });
+
 		const connection = await initTestDb(true);
 		Notes = connection.getRepository(MiNote);
 		alice = await signup({ username: 'alice' });

packages/backend/test/e2e/endpoints.ts

@@ -10,7 +10,7 @@ import * as assert from 'assert';
 // https://github.com/node-fetch/node-fetch/pull/1664
 import { Blob } from 'node-fetch';
 import { MiUser } from '@/models/_.js';
-import { api, initTestDb, post, signup, simpleGet, uploadFile } from '../utils.js';
+import { api, initTestDb, post, sendEnvUpdateRequest, signup, simpleGet, uploadFile } from '../utils.js';
 import type * as misskey from 'misskey-js';
 
 describe('Endpoints', () => {
@@ -20,6 +20,8 @@ describe('Endpoints', () => {
 	let dave: misskey.entities.SignupResponse;
 
 	beforeAll(async () => {
+		await sendEnvUpdateRequest({ key: 'FORCE_IGNORE_IDEMPOTENCY_FOR_TESTING', value: 'true' });
+
 		alice = await signup({ username: 'alice' });
 		bob = await signup({ username: 'bob' });
 		carol = await signup({ username: 'carol' });