キャッシュを導入してサーバーの負荷を下げた

cp-20 · cp-20 · commit 76780d7cb3a8 · 2024-12-26T23:09:06.000+09:00
diff --git a/apps/server/src/cache.test.ts b/apps/server/src/cache.test.ts
@@ -1,11 +1,11 @@
-import { memorize, TTLCache } from '@/cache';
+import { memorize, memorizeWithKey, purgeCache, TTLCache } from '@/cache';
 import { sleep } from 'bun';
 import { describe, test, setSystemTime, expect, mock } from 'bun:test';
 
 describe('TTLCache', () => {
   test('10sのTTLでキャッシュが期限切れになる', async () => {
     const fetcher = mock(async () => 'value');
-    const cache = new TTLCache<string, string>(10000);
+    const cache = new TTLCache<string, string>(10);
 
     setSystemTime(new Date('2024-01-01T00:00:00Z'));
     expect(await cache.get('key', fetcher), 'value').toBe('value');
@@ -21,7 +21,7 @@ describe('TTLCache', () => {
   });
 
   test('同一期間に二重でリクエストが飛ばない', async () => {
-    const cache = new TTLCache<string, string>(10000);
+    const cache = new TTLCache<string, string>(10);
 
     const fetcher = mock(async () => {
       await sleep(100);
@@ -42,7 +42,7 @@ describe('TTLCache', () => {
 describe('memorize', () => {
   test('10sのTTLでキャッシュが期限切れになる', async () => {
     const fetcher = mock(async (a: number, b: number) => a + b);
-    const memorized = memorize(10000, fetcher);
+    const memorized = memorize(10, fetcher);
 
     setSystemTime(new Date('2024-01-01T00:00:00Z'));
     expect(await memorized(1, 2)).toBe(3);
@@ -68,3 +68,35 @@ describe('memorize', () => {
     expect(fetcher).toHaveBeenCalledTimes(1);
   });
 });
+
+describe('memorizeWithKey', () => {
+  test('キャッシュされる', async () => {
+    const fetcher = mock(async (a: number, b: number) => a + b);
+    const memorized = memorizeWithKey(10, 'key', fetcher);
+
+    setSystemTime(new Date('2024-01-01T00:00:00Z'));
+    expect(await memorized(1, 2)).toBe(3);
+    expect(fetcher).toHaveBeenCalledTimes(1);
+
+    setSystemTime(new Date('2024-01-01T00:00:01Z'));
+    expect(await memorized(1, 2)).toBe(3);
+    expect(fetcher).toHaveBeenCalledTimes(1);
+  });
+
+  test('purge できる', async () => {
+    const fetcher = mock(async (a: number, b: number) => a + b);
+    const memorized = memorizeWithKey(10, 'key', fetcher);
+
+    setSystemTime(new Date('2024-01-01T00:00:00Z'));
+    expect(await memorized(1, 2)).toBe(3);
+    expect(fetcher).toHaveBeenCalledTimes(1);
+
+    setSystemTime(new Date('2024-01-01T00:00:01Z'));
+    expect(await memorized(1, 2)).toBe(3);
+    expect(fetcher).toHaveBeenCalledTimes(1);
+
+    purgeCache('key');
+    expect(await memorized(1, 2)).toBe(3);
+    expect(fetcher).toHaveBeenCalledTimes(2);
+  });
+});
diff --git a/apps/server/src/cache.ts b/apps/server/src/cache.ts
@@ -20,7 +20,7 @@ export class TTLCache<K, V> {
   private setCache(key: K, value: V) {
     this.cache.set(key, {
       value,
-      expiresAt: Date.now() + this.ttl,
+      expiresAt: Date.now() + this.ttl * 1000,
     });
   }
 
@@ -56,6 +56,15 @@ export class TTLCache<K, V> {
 
     return pendingPromise;
   }
+
+  list() {
+    return [...this.cache.keys()];
+  }
+
+  purge() {
+    // purge all cache
+    this.cache.clear();
+  }
 }
 
 export const memorize = <Args extends unknown[], Result>(
@@ -71,3 +80,29 @@ export const memorize = <Args extends unknown[], Result>(
 
   return fetcher;
 };
+
+const masterCache = new Map<string, TTLCache<string, unknown>>();
+
+export const getCaches = () => {
+  return [...masterCache].map(([key, cache]) => ({ key, cache: cache.list() }));
+};
+
+export const purgeCache = (key: string) => {
+  masterCache.get(key)?.purge();
+};
+
+export const memorizeWithKey = <Args extends unknown[], Result>(
+  ttl: number,
+  key: string,
+  fn: (...args: Args) => Promise<Result>,
+): ((...args: Args) => Promise<Result>) => {
+  const cache = new TTLCache<string, Result>(ttl);
+  masterCache.set(key, cache);
+
+  const fetcher = async (...args: Args): Promise<Result> => {
+    const key = JSON.stringify(args);
+    return cache.get(key, () => fn(...args));
+  };
+
+  return fetcher;
+};
diff --git a/apps/server/src/gateway.ts b/apps/server/src/gateway.ts
@@ -1,16 +1,33 @@
-import { memorize } from '@/cache';
+import { memorize, memorizeWithKey, purgeCache } from '@/cache';
 import { api, machineToken } from '@/traQ/api';
+import {
+  getChannelMessageRanking,
+  getChannelStampsRanking,
+  getGaveMessageStampsRanking,
+  getMessageContents,
+  getMessages,
+  getMessagesRanking,
+  getMessagesTimeline,
+  getReceivedMessageStampsRanking,
+  getStampRanking,
+  getStampRelations,
+  getStamps,
+  getStampsMeanUsage,
+  getSubscriptionRanking,
+  getTagRanking,
+  getUserGroupRanking,
+} from '@traq-ing/database';
 import sharp from 'sharp';
 
 const getAuthHeader = (token: string) => ({ headers: { Authorization: `Bearer ${token}` } });
 
-export const getMe = memorize(1000, async (token: string) => {
+export const getMe = memorize(10, async (token: string) => {
   const res = await api.users.getMe(getAuthHeader(token));
   if (!res.ok) throw new Error('Failed to fetch user');
   return res.data;
 });
 
-export const getSubscriptions = memorize(100, async (token: string) => {
+export const getSubscriptions = memorize(0, async (token: string) => {
   const res = await api.users.getMyChannelSubscriptions(getAuthHeader(token));
   if (!res.ok) throw new Error('Failed to fetch subscriptions');
   return res.data;
@@ -22,13 +39,13 @@ export const setSubscriptionLevel = async (token: string, channelId: string, lev
   return res.data;
 };
 
-export const getMessage = memorize(1000, async (messageId: string) => {
+export const getMessage = memorize(3600, async (messageId: string) => {
   const res = await api.messages.getMessage(messageId);
   if (!res.ok) throw new Error('Failed to fetch message');
   return res.data;
 });
 
-export const getUsers = memorize(1000, async () => {
+export const getUsers = memorize(3600, async () => {
   const res = await api.users.getUsers({ 'include-suspended': true });
   if (!res.ok) throw new Error('Failed to fetch users');
   return res.data;
@@ -103,3 +120,45 @@ export const getOgInfo = memorize(0, async (url: string) => {
 
   return { title, description, image, origin, type };
 });
+
+export const getMessagesCached = memorizeWithKey(3600, 'messages', getMessages);
+export const getMessageContentsCached = memorizeWithKey(3600, 'messageContents', getMessageContents);
+export const getChannelMessageRankingCached = memorizeWithKey(3600, 'channelMessageRanking', getChannelMessageRanking);
+export const getMessagesRankingCached = memorizeWithKey(3600, 'messagesRanking', getMessagesRanking);
+export const getMessagesTimelineCached = memorizeWithKey(3600, 'messagesTimeline', getMessagesTimeline);
+export const getStampRankingCached = memorizeWithKey(3600, 'stampRanking', getStampRanking);
+export const getChannelStampsRankingCached = memorizeWithKey(3600, 'channelStampsRanking', getChannelStampsRanking);
+export const getGaveMessageStampsRankingCached = memorizeWithKey(
+  3600,
+  'gaveMessageStampsRanking',
+  getGaveMessageStampsRanking,
+);
+export const getReceivedMessageStampsRankingCached = memorizeWithKey(
+  3600,
+  'receivedMessageStampsRanking',
+  getReceivedMessageStampsRanking,
+);
+export const getUserGroupRankingCached = memorizeWithKey(3600, 'userGroupRanking', getUserGroupRanking);
+export const getTagRankingCached = memorizeWithKey(3600, 'tagRanking', getTagRanking);
+export const getSubscriptionRankingCached = memorizeWithKey(3600, 'subscriptionRanking', getSubscriptionRanking);
+export const getStampsCached = memorizeWithKey(3600, 'stamps', getStamps);
+export const getStampsMeanUsageCached = memorizeWithKey(3600, 'stampsMeanUsage', getStampsMeanUsage);
+export const getStampRelationsCached = memorizeWithKey(3600, 'stampRelations', getStampRelations);
+
+export const forgotCaches = () => {
+  purgeCache('messages');
+  purgeCache('messageContents');
+  purgeCache('channelMessageRanking');
+  purgeCache('messagesRanking');
+  purgeCache('messagesTimeline');
+  purgeCache('stampRanking');
+  purgeCache('channelStampsRanking');
+  purgeCache('gaveMessageStampsRanking');
+  purgeCache('receivedMessageStampsRanking');
+  purgeCache('userGroupRanking');
+  purgeCache('tagRanking');
+  purgeCache('subscriptionRanking');
+  purgeCache('stamps');
+  purgeCache('stampsMeanUsage');
+  purgeCache('stampRelations');
+};
diff --git a/apps/server/src/index.ts b/apps/server/src/index.ts
@@ -4,43 +4,50 @@ import {
   MessagesQuerySchema,
   StampRelationsQuerySchema,
   StampsQuerySchema,
-  getChannelMessageRanking,
-  getChannelStampsRanking,
-  getGaveMessageStampsRanking,
-  getMessages,
-  getMessageContents,
-  getMessagesRanking,
-  getMessagesTimeline,
-  getReceivedMessageStampsRanking,
-  getStampRanking,
-  getStampRelations,
-  getStamps,
-  getSubscriptionRanking,
-  getTagRanking,
-  getUserGroupRanking,
   StampsMeanUsageQuerySchema,
-  getStampsMeanUsage,
 } from '@traq-ing/database';
 import { Hono } from 'hono';
 import { getCookie } from 'hono/cookie';
 import { z } from 'zod';
 import { tokenKey, traqAuthRoutes } from './auth';
 import {
+  forgotCaches,
+  getChannelMessageRankingCached,
   getChannels,
+  getChannelStampsRankingCached,
   getChannelSubscribers,
   getFile,
+  getGaveMessageStampsRankingCached,
   getMe,
   getMessage,
+  getMessageContentsCached,
+  getMessagesCached,
+  getMessagesRankingCached,
   getMessageStamps,
+  getMessagesTimelineCached,
   getOgInfo,
+  getReceivedMessageStampsRankingCached,
+  getStampRankingCached,
+  getStampRelationsCached,
+  getStampsCached,
+  getStampsMeanUsageCached,
+  getSubscriptionRankingCached,
   getSubscriptions,
+  getTagRankingCached,
+  getUserGroupRankingCached,
   getUserGroups,
   getUsers,
   setSubscriptionLevel,
 } from '@/gateway';
 import { HTTPException } from 'hono/http-exception';
 import { createMiddleware } from 'hono/factory';
 import { extractWords } from '@/extractor';
+import { getCaches } from '@/cache';
+
+const secretKey = process.env.SECRET_KEY;
+if (secretKey === undefined) {
+  throw new Error('SECRET_KEY is required');
+}
 
 const app = new Hono<{
   Variables: { token: string };
@@ -84,15 +91,15 @@ const routes = app
     },
   )
   .get('/messages', zValidator('query', MessagesQuerySchema), cacheMiddleware, async (c) =>
-    c.json(await getMessages(c.req.valid('query')), 200),
+    c.json(await getMessagesCached(c.req.valid('query')), 200),
   )
   .get(
     '/message-contents',
     zValidator('query', MessageContentsQuerySchema.omit({ offset: true })),
     cacheMiddleware,
     async (c) => {
       const { limit, ...query } = c.req.valid('query');
-      const contents = await getMessageContents(query);
+      const contents = await getMessageContentsCached(query);
       const words = [];
       for (const content of contents) {
         words.push(...(await extractWords(content.content)));
@@ -105,40 +112,42 @@ const routes = app
       return c.json(wordCount.slice(0, limit), 200);
     },
   )
-  .get('/channel-messages-ranking', cacheMiddleware, async (c) => c.json(await getChannelMessageRanking(), 200))
-  .get('/messages-ranking', cacheMiddleware, async (c) => c.json(await getMessagesRanking(), 200))
-  .get('/messages-timeline', cacheMiddleware, async (c) => c.json(await getMessagesTimeline(), 200))
-  .get('/stamp-ranking', cacheMiddleware, async (c) => c.json(await getStampRanking(), 200))
-  .get('/channel-stamps-ranking', cacheMiddleware, async (c) => c.json(await getChannelStampsRanking(), 200))
-  .get('/gave-stamps-ranking', cacheMiddleware, async (c) => c.json(await getGaveMessageStampsRanking(), 200))
-  .get('/received-stamps-ranking', cacheMiddleware, async (c) => c.json(await getReceivedMessageStampsRanking(), 200))
+  .get('/channel-messages-ranking', cacheMiddleware, async (c) => c.json(await getChannelMessageRankingCached(), 200))
+  .get('/messages-ranking', cacheMiddleware, async (c) => c.json(await getMessagesRankingCached(), 200))
+  .get('/messages-timeline', cacheMiddleware, async (c) => c.json(await getMessagesTimelineCached(), 200))
+  .get('/stamp-ranking', cacheMiddleware, async (c) => c.json(await getStampRankingCached(), 200))
+  .get('/channel-stamps-ranking', cacheMiddleware, async (c) => c.json(await getChannelStampsRankingCached(), 200))
+  .get('/gave-stamps-ranking', cacheMiddleware, async (c) => c.json(await getGaveMessageStampsRankingCached(), 200))
+  .get('/received-stamps-ranking', cacheMiddleware, async (c) =>
+    c.json(await getReceivedMessageStampsRankingCached(), 200),
+  )
   .get(
     '/group-ranking',
     zValidator('query', z.object({ groupBy: z.enum(['user', 'group']) })),
     cacheMiddleware,
-    async (c) => c.json(await getUserGroupRanking(c.req.valid('query').groupBy), 200),
+    async (c) => c.json(await getUserGroupRankingCached(c.req.valid('query').groupBy), 200),
   )
   .get(
     '/tag-ranking',
     zValidator('query', z.object({ groupBy: z.enum(['user', 'tag']) })),
     cacheMiddleware,
-    async (c) => c.json(await getTagRanking(c.req.valid('query').groupBy), 200),
+    async (c) => c.json(await getTagRankingCached(c.req.valid('query').groupBy), 200),
   )
   .get(
     '/subscription-ranking',
     zValidator('query', z.object({ groupBy: z.enum(['user', 'channel']) })),
     cacheMiddleware,
-    async (c) => c.json(await getSubscriptionRanking(c.req.valid('query').groupBy), 200),
+    async (c) => c.json(await getSubscriptionRankingCached(c.req.valid('query').groupBy), 200),
   )
   .get('/messages/:id', cacheMiddleware, async (c) => c.json(await getMessage(c.req.param('id')), 200))
   .get('/stamps', zValidator('query', StampsQuerySchema), async (c) =>
-    c.json(await getStamps(c.req.valid('query')), 200),
+    c.json(await getStampsCached(c.req.valid('query')), 200),
   )
   .get('/stamps-mean-usage', zValidator('query', StampsMeanUsageQuerySchema), async (c) =>
-    c.json(await getStampsMeanUsage(c.req.valid('query')), 200),
+    c.json(await getStampsMeanUsageCached(c.req.valid('query')), 200),
   )
   .get('/stamp-relations', zValidator('query', StampRelationsQuerySchema), async (c) =>
-    c.json(await getStampRelations(c.req.valid('query')), 200),
+    c.json(await getStampRelationsCached(c.req.valid('query')), 200),
   )
   .get('/users', cacheMiddleware, async (c) => c.json(await getUsers(), 200))
   .get('/groups', cacheMiddleware, async (c) => c.json(await getUserGroups(), 200))
@@ -170,7 +179,22 @@ const routes = app
       c.header('Cache-Control', 'private, max-age=31536000');
       return c.newResponse(file.stream());
     },
-  );
+  )
+  .get('/caches', async (c) => {
+    if (c.req.header('x-secret-key') !== secretKey) {
+      throw new HTTPException(401, { message: 'Unauthorized' });
+    }
+
+    return c.json(getCaches(), 200);
+  })
+  .delete('/caches', async (c) => {
+    if (c.req.header('x-secret-key') !== secretKey) {
+      throw new HTTPException(401, { message: 'Unauthorized' });
+    }
+    forgotCaches();
+    console.log('Caches cleared');
+    return c.json({ message: 'Caches cleared' }, 204);
+  });
 
 app.on(
   'GET',
diff --git a/apps/server/src/traQ/scripts/update-all-messages.ts b/apps/server/src/traQ/scripts/update-all-messages.ts
@@ -1,4 +1,10 @@
 import { dropMessagesAndStamps, updateMessages } from '@/traQ';
 
+const baseUrl = process.env.SERVER_URL;
+if (!baseUrl) throw new Error('SERVER_URL is not set');
+const secret = process.env.SECRET_KEY;
+if (!secret) throw new Error('SECRET_KEY is not set');
+
 await dropMessagesAndStamps();
 await updateMessages(true);
+await fetch(`${baseUrl}/caches`, { method: 'DELETE', headers: { 'X-Secret-Key': secret } });
diff --git a/apps/server/src/traQ/scripts/update-messages.ts b/apps/server/src/traQ/scripts/update-messages.ts
diff --git a/apps/server/src/traQ/scripts/update-statistics.ts b/apps/server/src/traQ/scripts/update-statistics.ts
