From 7ec07d5fd2e28d523c85d160944916b2a5ee04a1 Mon Sep 17 00:00:00 2001
From: tamaina <tamaina@hotmail.co.jp>
Date: Sat, 8 Jul 2023 21:18:16 +0900
Subject: perf(backend): Reduce memory usage of MemoryKVCache (#11076)

* perf(backend): Reduce memory usage of MemoryKVCache

* fix
---
 packages/backend/src/core/CacheService.ts | 47 +++++++++++++++++++++++++------
 1 file changed, 39 insertions(+), 8 deletions(-)

(limited to 'packages/backend/src/core')

diff --git a/packages/backend/src/core/CacheService.ts b/packages/backend/src/core/CacheService.ts
index 2b7f9a48da..cd6b68e721 100644
--- a/packages/backend/src/core/CacheService.ts
+++ b/packages/backend/src/core/CacheService.ts
@@ -11,10 +11,10 @@ import type { OnApplicationShutdown } from '@nestjs/common';
 
 @Injectable()
 export class CacheService implements OnApplicationShutdown {
-	public userByIdCache: MemoryKVCache<User>;
-	public localUserByNativeTokenCache: MemoryKVCache<LocalUser | null>;
+	public userByIdCache: MemoryKVCache<User, User | string>;
+	public localUserByNativeTokenCache: MemoryKVCache<LocalUser | null, string | null>;
 	public localUserByIdCache: MemoryKVCache<LocalUser>;
-	public uriPersonCache: MemoryKVCache<User | null>;
+	public uriPersonCache: MemoryKVCache<User | null, string | null>;
 	public userProfileCache: RedisKVCache<UserProfile>;
 	public userMutingsCache: RedisKVCache<Set<string>>;
 	public userBlockingCache: RedisKVCache<Set<string>>;
@@ -55,10 +55,41 @@ export class CacheService implements OnApplicationShutdown {
 	) {
 		//this.onMessage = this.onMessage.bind(this);
 
-		this.userByIdCache = new MemoryKVCache<User>(Infinity);
-		this.localUserByNativeTokenCache = new MemoryKVCache<LocalUser | null>(Infinity);
-		this.localUserByIdCache = new MemoryKVCache<LocalUser>(Infinity);
-		this.uriPersonCache = new MemoryKVCache<User | null>(Infinity);
+		const localUserByIdCache = new MemoryKVCache<LocalUser>(1000 * 60 * 60 * 6 /* 6h */);
+		this.localUserByIdCache	= localUserByIdCache;
+
+		// ローカルユーザーならlocalUserByIdCacheにデータを追加し、こちらにはid(文字列)だけを追加する
+		const userByIdCache = new MemoryKVCache<User, User | string>(1000 * 60 * 60 * 6 /* 6h */, {
+			toMapConverter: user => {
+				if (user.host === null) {
+					localUserByIdCache.set(user.id, user as LocalUser);
+					return user.id;
+				}
+
+				return user;
+			},
+			fromMapConverter: userOrId => typeof userOrId === 'string' ? localUserByIdCache.get(userOrId) : userOrId,
+		});
+		this.userByIdCache = userByIdCache;
+
+		this.localUserByNativeTokenCache = new MemoryKVCache<LocalUser | null, string | null>(Infinity, {
+			toMapConverter: user => {
+				if (user === null) return null;
+
+				localUserByIdCache.set(user.id, user);
+				return user.id;
+			},
+			fromMapConverter: id => id === null ? null : localUserByIdCache.get(id),
+		});
+		this.uriPersonCache = new MemoryKVCache<User | null, string | null>(Infinity, {
+			toMapConverter: user => {
+				if (user === null) return null;
+
+				userByIdCache.set(user.id, user);
+				return user.id;
+			},
+			fromMapConverter: id => id === null ? null : userByIdCache.get(id),
+		});
 
 		this.userProfileCache = new RedisKVCache<UserProfile>(this.redisClient, 'userProfile', {
 			lifetime: 1000 * 60 * 30, // 30m
@@ -131,7 +162,7 @@ export class CacheService implements OnApplicationShutdown {
 					const user = await this.usersRepository.findOneByOrFail({ id: body.id });
 					this.userByIdCache.set(user.id, user);
 					for (const [k, v] of this.uriPersonCache.cache.entries()) {
-						if (v.value?.id === user.id) {
+						if (v.value === user.id) {
 							this.uriPersonCache.set(k, user);
 						}
 					}
-- 
cgit v1.2.3-freya