From 9253e522aa7c04008b3a18a51ca0d988704d0877 Mon Sep 17 00:00:00 2001 From: liuyi Date: Thu, 11 Jan 2024 06:40:53 +0000 Subject: [PATCH] test(server): avoid progress get hold after tests finished (#5522) --- packages/backend/server/src/affine.config.ts | 38 +++++++----- packages/backend/server/src/affine.env.ts | 12 ++++ packages/backend/server/src/app.ts | 12 +++- packages/backend/server/src/cache/index.ts | 40 ++++++++----- .../backend/server/src/cache/instances.ts | 4 ++ .../backend/server/src/cache/interceptor.ts | 2 +- .../server/src/cache/{ => providers}/cache.ts | 6 +- .../server/src/cache/{ => providers}/redis.ts | 0 .../backend/server/src/cache/redis/index.ts | 38 ++++++++++++ packages/backend/server/src/config/def.ts | 4 ++ packages/backend/server/src/config/default.ts | 3 + packages/backend/server/src/config/index.ts | 2 - packages/backend/server/src/data/app.ts | 5 ++ .../server/src/data/commands/create.ts | 6 +- .../src/data/migrations/1698398506533-guid.ts | 4 +- .../1698652531198-user-features-init.ts | 7 ++- .../1699005339766-page-permission.ts | 7 +-- .../1702620653283-old-user-feature.ts | 8 +-- .../1703756315970-unamed-account.ts | 7 ++- .../1703828796699-workspace-blobs.ts | 6 +- .../1704352562369-refresh-user-features.ts | 7 ++- .../data/migrations/utils/user-features.ts | 7 +-- packages/backend/server/src/index.ts | 17 +++--- packages/backend/server/src/metrics/index.ts | 32 +++++++++- .../backend/server/src/metrics/metrics.ts | 2 + .../server/src/metrics/opentelemetry.ts | 28 +-------- .../server/src/modules/auth/resolver.ts | 7 ++- .../server/src/modules/sync/redis-adapter.ts | 14 +---- packages/backend/server/src/prelude.ts | 14 +---- packages/backend/server/src/session.ts | 47 +++------------ packages/backend/server/src/throttler.ts | 60 +++++++++++-------- packages/backend/server/tests/auth.spec.ts | 2 + packages/backend/server/tests/feature.spec.ts | 2 + packages/backend/server/tests/mailer.e2e.ts | 2 + packages/backend/server/tests/quota.spec.ts | 2 + packages/backend/server/tests/session.spec.ts | 2 + 36 files changed, 262 insertions(+), 194 deletions(-) create mode 100644 packages/backend/server/src/affine.env.ts create mode 100644 packages/backend/server/src/cache/instances.ts rename packages/backend/server/src/cache/{ => providers}/cache.ts (98%) rename packages/backend/server/src/cache/{ => providers}/redis.ts (100%) create mode 100644 packages/backend/server/src/cache/redis/index.ts diff --git a/packages/backend/server/src/affine.config.ts b/packages/backend/server/src/affine.config.ts index 43e9a83fac..869b20467d 100644 --- a/packages/backend/server/src/affine.config.ts +++ b/packages/backend/server/src/affine.config.ts @@ -5,21 +5,27 @@ const env = process.env; const node = AFFiNE.node; // TODO: may be separate config overring in `affine.[env].config`? -if (node.prod && env.R2_OBJECT_STORAGE_ACCOUNT_ID) { - AFFiNE.storage.providers.r2 = { - accountId: env.R2_OBJECT_STORAGE_ACCOUNT_ID, - credentials: { - accessKeyId: env.R2_OBJECT_STORAGE_ACCESS_KEY_ID!, - secretAccessKey: env.R2_OBJECT_STORAGE_SECRET_ACCESS_KEY!, - }, - }; - AFFiNE.storage.storages.avatar.provider = 'r2'; - AFFiNE.storage.storages.avatar.bucket = 'account-avatar'; - AFFiNE.storage.storages.avatar.publicLinkFactory = key => - `https://avatar.affineassets.com/${key}`; +if (node.prod) { + // Storage + if (env.R2_OBJECT_STORAGE_ACCOUNT_ID) { + AFFiNE.storage.providers.r2 = { + accountId: env.R2_OBJECT_STORAGE_ACCOUNT_ID, + credentials: { + accessKeyId: env.R2_OBJECT_STORAGE_ACCESS_KEY_ID!, + secretAccessKey: env.R2_OBJECT_STORAGE_SECRET_ACCESS_KEY!, + }, + }; + AFFiNE.storage.storages.avatar.provider = 'r2'; + AFFiNE.storage.storages.avatar.bucket = 'account-avatar'; + AFFiNE.storage.storages.avatar.publicLinkFactory = key => + `https://avatar.affineassets.com/${key}`; - AFFiNE.storage.storages.blob.provider = 'r2'; - AFFiNE.storage.storages.blob.bucket = `workspace-blobs-${ - AFFiNE.affine.canary ? 'canary' : 'prod' - }`; + AFFiNE.storage.storages.blob.provider = 'r2'; + AFFiNE.storage.storages.blob.bucket = `workspace-blobs-${ + AFFiNE.affine.canary ? 'canary' : 'prod' + }`; + } + + // Metrics + AFFiNE.metrics.enabled = true; } diff --git a/packages/backend/server/src/affine.env.ts b/packages/backend/server/src/affine.env.ts new file mode 100644 index 0000000000..5605964bf5 --- /dev/null +++ b/packages/backend/server/src/affine.env.ts @@ -0,0 +1,12 @@ +import { homedir } from 'node:os'; +import { join } from 'node:path'; + +import { config } from 'dotenv'; + +import { SERVER_FLAVOR } from './config/default'; + +if (SERVER_FLAVOR === 'selfhosted') { + config({ path: join(homedir(), '.affine', '.env') }); +} else { + config(); +} diff --git a/packages/backend/server/src/app.ts b/packages/backend/server/src/app.ts index ad97741aac..b123026318 100644 --- a/packages/backend/server/src/app.ts +++ b/packages/backend/server/src/app.ts @@ -3,24 +3,32 @@ import { APP_INTERCEPTOR } from '@nestjs/core'; import { AppController } from './app.controller'; import { CacheInterceptor, CacheModule } from './cache'; +import { RedisModule } from './cache/redis'; import { ConfigModule, SERVER_FLAVOR } from './config'; import { EventModule } from './event'; +import { MetricsModule } from './metrics'; import { BusinessModules } from './modules'; import { AuthModule } from './modules/auth'; import { PrismaModule } from './prisma'; import { SessionModule } from './session'; import { RateLimiterModule } from './throttler'; -const BasicModules = [ - PrismaModule, +export const BasicModules = [ ConfigModule.forRoot(), CacheModule, + PrismaModule, + MetricsModule, EventModule, SessionModule, RateLimiterModule, AuthModule, ]; +// better module registration logic +if (AFFiNE.redis.enabled) { + BasicModules.push(RedisModule); +} + @Module({ providers: [ { diff --git a/packages/backend/server/src/cache/index.ts b/packages/backend/server/src/cache/index.ts index 3a1dab5310..06c73367f2 100644 --- a/packages/backend/server/src/cache/index.ts +++ b/packages/backend/server/src/cache/index.ts @@ -1,26 +1,34 @@ -import { FactoryProvider, Global, Module } from '@nestjs/common'; +import { Global, Module, Provider, Type } from '@nestjs/common'; import { Redis } from 'ioredis'; -import { Config } from '../config'; -import { LocalCache } from './cache'; -import { RedisCache } from './redis'; +import { SessionCache, ThrottlerCache } from './instances'; +import { LocalCache } from './providers/cache'; +import { RedisCache } from './providers/redis'; +import { CacheRedis, SessionRedis, ThrottlerRedis } from './redis'; -const CacheProvider: FactoryProvider = { - provide: LocalCache, - useFactory: (config: Config) => { - return config.redis.enabled - ? new RedisCache(new Redis(config.redis)) - : new LocalCache(); - }, - inject: [Config], -}; +function makeCacheProvider(CacheToken: Type, RedisToken: Type): Provider { + return { + provide: CacheToken, + useFactory: (redis?: Redis) => { + return redis ? new RedisCache(redis) : new LocalCache(); + }, + inject: [{ token: RedisToken, optional: true }], + }; +} + +const CacheProvider = makeCacheProvider(LocalCache, CacheRedis); +const SessionCacheProvider = makeCacheProvider(SessionCache, SessionRedis); +const ThrottlerCacheProvider = makeCacheProvider( + ThrottlerCache, + ThrottlerRedis +); @Global() @Module({ - providers: [CacheProvider], - exports: [CacheProvider], + providers: [CacheProvider, SessionCacheProvider, ThrottlerCacheProvider], + exports: [CacheProvider, SessionCacheProvider, ThrottlerCacheProvider], }) export class CacheModule {} -export { LocalCache as Cache }; +export { LocalCache as Cache, SessionCache, ThrottlerCache }; export { CacheInterceptor, MakeCache, PreventCache } from './interceptor'; diff --git a/packages/backend/server/src/cache/instances.ts b/packages/backend/server/src/cache/instances.ts new file mode 100644 index 0000000000..0062aee8cc --- /dev/null +++ b/packages/backend/server/src/cache/instances.ts @@ -0,0 +1,4 @@ +import { LocalCache } from './providers/cache'; + +export class SessionCache extends LocalCache {} +export class ThrottlerCache extends LocalCache {} diff --git a/packages/backend/server/src/cache/interceptor.ts b/packages/backend/server/src/cache/interceptor.ts index 8cf078e18f..93cd388273 100644 --- a/packages/backend/server/src/cache/interceptor.ts +++ b/packages/backend/server/src/cache/interceptor.ts @@ -10,7 +10,7 @@ import { Reflector } from '@nestjs/core'; import { GqlContextType, GqlExecutionContext } from '@nestjs/graphql'; import { mergeMap, Observable, of } from 'rxjs'; -import { LocalCache } from './cache'; +import { LocalCache } from './providers/cache'; export const MakeCache = (key: string[], args?: string[]) => SetMetadata('cacheKey', [key, args]); diff --git a/packages/backend/server/src/cache/cache.ts b/packages/backend/server/src/cache/providers/cache.ts similarity index 98% rename from packages/backend/server/src/cache/cache.ts rename to packages/backend/server/src/cache/providers/cache.ts index 24df6ffe76..b770473bfb 100644 --- a/packages/backend/server/src/cache/cache.ts +++ b/packages/backend/server/src/cache/providers/cache.ts @@ -1,3 +1,4 @@ +import { Injectable } from '@nestjs/common'; import Keyv from 'keyv'; export interface CacheSetOptions { @@ -50,11 +51,12 @@ export interface Cache { mapLen(map: string): Promise; } +@Injectable() export class LocalCache implements Cache { private readonly kv: Keyv; - constructor() { - this.kv = new Keyv(); + constructor(opts: Keyv.Options = {}) { + this.kv = new Keyv(opts); } // standard operation diff --git a/packages/backend/server/src/cache/redis.ts b/packages/backend/server/src/cache/providers/redis.ts similarity index 100% rename from packages/backend/server/src/cache/redis.ts rename to packages/backend/server/src/cache/providers/redis.ts diff --git a/packages/backend/server/src/cache/redis/index.ts b/packages/backend/server/src/cache/redis/index.ts new file mode 100644 index 0000000000..d2807540c3 --- /dev/null +++ b/packages/backend/server/src/cache/redis/index.ts @@ -0,0 +1,38 @@ +import { Global, Injectable, Module, OnModuleDestroy } from '@nestjs/common'; +import { Redis as IORedis } from 'ioredis'; + +import { Config } from '../../config'; + +class Redis extends IORedis implements OnModuleDestroy { + onModuleDestroy() { + this.disconnect(); + } +} + +@Injectable() +export class CacheRedis extends Redis { + constructor(config: Config) { + super({ ...config.redis, db: config.redis.database }); + } +} + +@Injectable() +export class ThrottlerRedis extends Redis { + constructor(config: Config) { + super({ ...config.redis, db: config.redis.database + 1 }); + } +} + +@Injectable() +export class SessionRedis extends Redis { + constructor(config: Config) { + super({ ...config.redis, db: config.redis.database + 2 }); + } +} + +@Global() +@Module({ + providers: [CacheRedis, ThrottlerRedis, SessionRedis], + exports: [CacheRedis, ThrottlerRedis, SessionRedis], +}) +export class RedisModule {} diff --git a/packages/backend/server/src/config/def.ts b/packages/backend/server/src/config/def.ts index 34221ff3d9..ceb871ee1f 100644 --- a/packages/backend/server/src/config/def.ts +++ b/packages/backend/server/src/config/def.ts @@ -357,6 +357,10 @@ export interface AFFiNEConfig { }; }; + metrics: { + enabled: boolean; + }; + payment: { stripe: { keys: { diff --git a/packages/backend/server/src/config/default.ts b/packages/backend/server/src/config/default.ts index eca97146f2..94c8e72793 100644 --- a/packages/backend/server/src/config/default.ts +++ b/packages/backend/server/src/config/default.ts @@ -211,6 +211,9 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => { apiVersion: '2023-10-16', }, }, + metrics: { + enabled: false, + }, } satisfies AFFiNEConfig; applyEnvToConfig(defaultConfig); diff --git a/packages/backend/server/src/config/index.ts b/packages/backend/server/src/config/index.ts index 722dee3e93..433221764e 100644 --- a/packages/backend/server/src/config/index.ts +++ b/packages/backend/server/src/config/index.ts @@ -5,8 +5,6 @@ import { merge } from 'lodash-es'; import type { DeepPartial } from '../utils/types'; import type { AFFiNEConfig } from './def'; -import '../prelude'; - type ConstructorOf = { new (): T; }; diff --git a/packages/backend/server/src/data/app.ts b/packages/backend/server/src/data/app.ts index 5287e74fd9..5e00fd214a 100644 --- a/packages/backend/server/src/data/app.ts +++ b/packages/backend/server/src/data/app.ts @@ -1,3 +1,5 @@ +import '../prelude'; + import { Logger, Module } from '@nestjs/common'; import { CommandFactory } from 'nest-commander'; @@ -14,6 +16,9 @@ import { RevertCommand, RunCommand } from './commands/run'; enableUpdateAutoMerging: false, }, }, + metrics: { + enabled: false, + }, }), BusinessAppModule, ], diff --git a/packages/backend/server/src/data/commands/create.ts b/packages/backend/server/src/data/commands/create.ts index f8faf3a744..dba829aa2d 100644 --- a/packages/backend/server/src/data/commands/create.ts +++ b/packages/backend/server/src/data/commands/create.ts @@ -59,13 +59,13 @@ export class CreateCommand extends CommandRunner { } private createScript(name: string) { - const contents = ["import { PrismaService } from '../../prisma';", '']; + const contents = ["import { PrismaClient } from '@prisma/client';", '']; contents.push(`export class ${name} {`); contents.push(' // do the migration'); - contents.push(' static async up(db: PrismaService) {}'); + contents.push(' static async up(db: PrismaClient) {}'); contents.push(''); contents.push(' // revert the migration'); - contents.push(' static async down(db: PrismaService) {}'); + contents.push(' static async down(db: PrismaClient) {}'); contents.push('}'); diff --git a/packages/backend/server/src/data/migrations/1698398506533-guid.ts b/packages/backend/server/src/data/migrations/1698398506533-guid.ts index 4ace563ab1..e7e0262472 100644 --- a/packages/backend/server/src/data/migrations/1698398506533-guid.ts +++ b/packages/backend/server/src/data/migrations/1698398506533-guid.ts @@ -1,11 +1,11 @@ +import { PrismaClient } from '@prisma/client'; import { applyUpdate, Doc, encodeStateAsUpdate } from 'yjs'; -import { PrismaService } from '../../prisma'; import { DocID } from '../../utils/doc'; export class Guid1698398506533 { // do the migration - static async up(db: PrismaService) { + static async up(db: PrismaClient) { let turn = 0; let lastTurnCount = 100; while (lastTurnCount === 100) { diff --git a/packages/backend/server/src/data/migrations/1698652531198-user-features-init.ts b/packages/backend/server/src/data/migrations/1698652531198-user-features-init.ts index 786d467718..7d28a18d40 100644 --- a/packages/backend/server/src/data/migrations/1698652531198-user-features-init.ts +++ b/packages/backend/server/src/data/migrations/1698652531198-user-features-init.ts @@ -1,11 +1,12 @@ +import { PrismaClient } from '@prisma/client'; + import { Features } from '../../modules/features'; import { Quotas } from '../../modules/quota/schema'; -import { PrismaService } from '../../prisma'; import { migrateNewFeatureTable, upsertFeature } from './utils/user-features'; export class UserFeaturesInit1698652531198 { // do the migration - static async up(db: PrismaService) { + static async up(db: PrismaClient) { // upgrade features from lower version to higher version for (const feature of Features) { await upsertFeature(db, feature); @@ -18,7 +19,7 @@ export class UserFeaturesInit1698652531198 { } // revert the migration - static async down(_db: PrismaService) { + static async down(_db: PrismaClient) { // TODO: revert the migration } } diff --git a/packages/backend/server/src/data/migrations/1699005339766-page-permission.ts b/packages/backend/server/src/data/migrations/1699005339766-page-permission.ts index ae8c5620fd..ca3238b335 100644 --- a/packages/backend/server/src/data/migrations/1699005339766-page-permission.ts +++ b/packages/backend/server/src/data/migrations/1699005339766-page-permission.ts @@ -1,8 +1,7 @@ -import { PrismaService } from '../../prisma'; - +import { PrismaClient } from '@prisma/client'; export class PagePermission1699005339766 { // do the migration - static async up(db: PrismaService) { + static async up(db: PrismaClient) { let turn = 0; let lastTurnCount = 50; const done = new Set(); @@ -88,7 +87,7 @@ export class PagePermission1699005339766 { } // revert the migration - static async down(db: PrismaService) { + static async down(db: PrismaClient) { await db.workspaceUserPermission.deleteMany({}); await db.workspacePageUserPermission.deleteMany({}); } diff --git a/packages/backend/server/src/data/migrations/1702620653283-old-user-feature.ts b/packages/backend/server/src/data/migrations/1702620653283-old-user-feature.ts index 330b33f433..2656ab814c 100644 --- a/packages/backend/server/src/data/migrations/1702620653283-old-user-feature.ts +++ b/packages/backend/server/src/data/migrations/1702620653283-old-user-feature.ts @@ -1,9 +1,9 @@ -import { QuotaType } from '../../modules/quota/types'; -import { PrismaService } from '../../prisma'; +import { PrismaClient } from '@prisma/client'; +import { QuotaType } from '../../modules/quota/types'; export class OldUserFeature1702620653283 { // do the migration - static async up(db: PrismaService) { + static async up(db: PrismaClient) { await db.$transaction(async tx => { const latestFreePlan = await tx.features.findFirstOrThrow({ where: { feature: QuotaType.FreePlanV1 }, @@ -31,7 +31,7 @@ export class OldUserFeature1702620653283 { // revert the migration // WARN: this will drop all user features - static async down(db: PrismaService) { + static async down(db: PrismaClient) { await db.userFeatures.deleteMany({}); } } diff --git a/packages/backend/server/src/data/migrations/1703756315970-unamed-account.ts b/packages/backend/server/src/data/migrations/1703756315970-unamed-account.ts index bd8b7b8d4c..1a4bfae491 100644 --- a/packages/backend/server/src/data/migrations/1703756315970-unamed-account.ts +++ b/packages/backend/server/src/data/migrations/1703756315970-unamed-account.ts @@ -1,9 +1,10 @@ +import { PrismaClient } from '@prisma/client'; + import type { UserType } from '../../modules/users'; -import { PrismaService } from '../../prisma'; export class UnamedAccount1703756315970 { // do the migration - static async up(db: PrismaService) { + static async up(db: PrismaClient) { await db.$transaction(async tx => { // only find users with empty names const users = await db.$queryRaw< @@ -27,5 +28,5 @@ export class UnamedAccount1703756315970 { } // revert the migration - static async down(_db: PrismaService) {} + static async down(_db: PrismaClient) {} } diff --git a/packages/backend/server/src/data/migrations/1703828796699-workspace-blobs.ts b/packages/backend/server/src/data/migrations/1703828796699-workspace-blobs.ts index 6bbeff1c3e..0e505cc1db 100644 --- a/packages/backend/server/src/data/migrations/1703828796699-workspace-blobs.ts +++ b/packages/backend/server/src/data/migrations/1703828796699-workspace-blobs.ts @@ -1,11 +1,11 @@ import { ModuleRef } from '@nestjs/core'; +import { PrismaClient } from '@prisma/client'; import { WorkspaceBlobStorage } from '../../modules/storage'; -import { PrismaService } from '../../prisma'; export class WorkspaceBlobs1703828796699 { // do the migration - static async up(db: PrismaService, injector: ModuleRef) { + static async up(db: PrismaClient, injector: ModuleRef) { const blobStorage = injector.get(WorkspaceBlobStorage, { strict: false }); let hasMore = true; let turn = 0; @@ -32,7 +32,7 @@ export class WorkspaceBlobs1703828796699 { } // revert the migration - static async down(_db: PrismaService) { + static async down(_db: PrismaClient) { // old data kept, no need to downgrade the migration } } diff --git a/packages/backend/server/src/data/migrations/1704352562369-refresh-user-features.ts b/packages/backend/server/src/data/migrations/1704352562369-refresh-user-features.ts index 407aa22986..1e903344ae 100644 --- a/packages/backend/server/src/data/migrations/1704352562369-refresh-user-features.ts +++ b/packages/backend/server/src/data/migrations/1704352562369-refresh-user-features.ts @@ -1,10 +1,11 @@ +import { PrismaClient } from '@prisma/client'; + import { Features } from '../../modules/features'; -import { PrismaService } from '../../prisma'; import { upsertFeature } from './utils/user-features'; export class RefreshUserFeatures1704352562369 { // do the migration - static async up(db: PrismaService) { + static async up(db: PrismaClient) { // add early access v2 & copilot feature for (const feature of Features) { await upsertFeature(db, feature); @@ -12,5 +13,5 @@ export class RefreshUserFeatures1704352562369 { } // revert the migration - static async down(_db: PrismaService) {} + static async down(_db: PrismaClient) {} } diff --git a/packages/backend/server/src/data/migrations/utils/user-features.ts b/packages/backend/server/src/data/migrations/utils/user-features.ts index 0e93ce2328..3483420771 100644 --- a/packages/backend/server/src/data/migrations/utils/user-features.ts +++ b/packages/backend/server/src/data/migrations/utils/user-features.ts @@ -1,15 +1,14 @@ -import { Prisma } from '@prisma/client'; +import { Prisma, PrismaClient } from '@prisma/client'; import { CommonFeature, FeatureKind, FeatureType, } from '../../../modules/features'; -import { PrismaService } from '../../../prisma'; // upgrade features from lower version to higher version export async function upsertFeature( - db: PrismaService, + db: PrismaClient, feature: CommonFeature ): Promise { const hasEqualOrGreaterVersion = @@ -34,7 +33,7 @@ export async function upsertFeature( } } -export async function migrateNewFeatureTable(prisma: PrismaService) { +export async function migrateNewFeatureTable(prisma: PrismaClient) { const waitingList = await prisma.newFeaturesWaitingList.findMany(); for (const oldUser of waitingList) { const user = await prisma.user.findFirst({ diff --git a/packages/backend/server/src/index.ts b/packages/backend/server/src/index.ts index 187dd1f920..6333e1798d 100644 --- a/packages/backend/server/src/index.ts +++ b/packages/backend/server/src/index.ts @@ -1,6 +1,7 @@ /// -import { start as startAutoMetrics } from './metrics'; -startAutoMetrics(); +// keep the config import at the top +// eslint-disable-next-line simple-import-sort/imports +import './prelude'; import { NestFactory } from '@nestjs/core'; import type { NestExpressApplication } from '@nestjs/platform-express'; @@ -12,6 +13,7 @@ import { Config } from './config'; import { ExceptionLogger } from './middleware/exception-logger'; import { serverTimingAndCache } from './middleware/timing'; import { RedisIoAdapter } from './modules/sync/redis-adapter'; +import { CacheRedis } from './cache/redis'; const { NODE_ENV, AFFINE_ENV } = process.env; const app = await NestFactory.create(AppModule, { @@ -42,15 +44,10 @@ const config = app.get(Config); const host = config.node.prod ? '0.0.0.0' : 'localhost'; const port = config.port ?? 3010; -if (config.redis.enabled) { +if (!config.node.test && config.redis.enabled) { + const redis = app.get(CacheRedis, { strict: false }); const redisIoAdapter = new RedisIoAdapter(app); - await redisIoAdapter.connectToRedis( - config.redis.host, - config.redis.port, - config.redis.username, - config.redis.password, - config.redis.database - ); + await redisIoAdapter.connectToRedis(redis); app.useWebSocketAdapter(redisIoAdapter); } diff --git a/packages/backend/server/src/metrics/index.ts b/packages/backend/server/src/metrics/index.ts index 655c717c3e..9ca470b582 100644 --- a/packages/backend/server/src/metrics/index.ts +++ b/packages/backend/server/src/metrics/index.ts @@ -1,3 +1,33 @@ +import { Global, Module, OnModuleDestroy, OnModuleInit } from '@nestjs/common'; +import { NodeSDK } from '@opentelemetry/sdk-node'; + +import { Config } from '../config'; +import { parseEnvValue } from '../config/def'; +import { createSDK, registerCustomMetrics } from './opentelemetry'; + +@Global() +@Module({}) +export class MetricsModule implements OnModuleInit, OnModuleDestroy { + private sdk: NodeSDK | null = null; + constructor(private readonly config: Config) {} + + onModuleInit() { + if ( + this.config.metrics.enabled && + !parseEnvValue(process.env.DISABLE_TELEMETRY, 'boolean') + ) { + this.sdk = createSDK(); + this.sdk.start(); + registerCustomMetrics(); + } + } + + async onModuleDestroy() { + if (this.config.metrics.enabled && this.sdk) { + await this.sdk.shutdown(); + } + } +} + export * from './metrics'; -export { start } from './opentelemetry'; export * from './utils'; diff --git a/packages/backend/server/src/metrics/metrics.ts b/packages/backend/server/src/metrics/metrics.ts index ccbb7be606..83e6d72216 100644 --- a/packages/backend/server/src/metrics/metrics.ts +++ b/packages/backend/server/src/metrics/metrics.ts @@ -127,3 +127,5 @@ export const metrics = new Proxy>( }, } ); + +export function stopMetrics() {} diff --git a/packages/backend/server/src/metrics/opentelemetry.ts b/packages/backend/server/src/metrics/opentelemetry.ts index 54e449ae13..0f3206bb2c 100644 --- a/packages/backend/server/src/metrics/opentelemetry.ts +++ b/packages/backend/server/src/metrics/opentelemetry.ts @@ -33,7 +33,6 @@ import prismaInstrument from '@prisma/instrumentation'; const { PrismaInstrumentation } = prismaInstrument; -import { parseEnvValue } from '../config/def'; import { PrismaMetricProducer } from './prisma'; abstract class OpentelemetryFactor { @@ -116,7 +115,8 @@ class DebugOpentelemetryFactor extends OpentelemetryFactor { } } -function createSDK() { +// TODO(@forehalo): make it configurable +export function createSDK() { let factor: OpentelemetryFactor | null = null; if (process.env.NODE_ENV === 'production') { factor = new GCloudOpentelemetryFactor(); @@ -129,21 +129,11 @@ function createSDK() { return factor?.create(); } -let OPENTELEMETRY_STARTED = false; - -function ensureStarted() { - if (!OPENTELEMETRY_STARTED) { - OPENTELEMETRY_STARTED = true; - start(); - } -} - function getMeterProvider() { - ensureStarted(); return metrics.getMeterProvider(); } -function registerCustomMetrics() { +export function registerCustomMetrics() { const hostMetricsMonitoring = new HostMetrics({ name: 'instance-host-metrics', meterProvider: getMeterProvider() as MeterProvider, @@ -154,15 +144,3 @@ function registerCustomMetrics() { export function getMeter(name = 'business') { return getMeterProvider().getMeter(name); } - -export function start() { - if (parseEnvValue(process.env.DISABLE_TELEMETRY, 'boolean')) { - return; - } - const sdk = createSDK(); - - if (sdk) { - sdk.start(); - registerCustomMetrics(); - } -} diff --git a/packages/backend/server/src/modules/auth/resolver.ts b/packages/backend/server/src/modules/auth/resolver.ts index 697812520f..aac7b4e673 100644 --- a/packages/backend/server/src/modules/auth/resolver.ts +++ b/packages/backend/server/src/modules/auth/resolver.ts @@ -167,8 +167,13 @@ export class AuthResolver { @CurrentUser() user: UserType, @Args('token') token: string ) { + const key = await this.session.get(token); + if (!key) { + throw new ForbiddenException('Invalid token'); + } + // email has set token in `sendVerifyChangeEmail` - const [id, email] = (await this.session.get(token)).split(','); + const [id, email] = key.split(','); if (!id || id !== user.id || !email) { throw new ForbiddenException('Invalid token'); } diff --git a/packages/backend/server/src/modules/sync/redis-adapter.ts b/packages/backend/server/src/modules/sync/redis-adapter.ts index 312a0c1d7e..46d1591975 100644 --- a/packages/backend/server/src/modules/sync/redis-adapter.ts +++ b/packages/backend/server/src/modules/sync/redis-adapter.ts @@ -6,18 +6,8 @@ import { ServerOptions } from 'socket.io'; export class RedisIoAdapter extends IoAdapter { private adapterConstructor: ReturnType | undefined; - async connectToRedis( - host: string, - port: number, - username: string, - password: string, - db: number - ): Promise { - const pubClient = new Redis(port, host, { - username, - password, - db, - }); + async connectToRedis(redis: Redis): Promise { + const pubClient = redis; pubClient.on('error', err => { console.error(err); }); diff --git a/packages/backend/server/src/prelude.ts b/packages/backend/server/src/prelude.ts index 2ca5041763..d70e339c2c 100644 --- a/packages/backend/server/src/prelude.ts +++ b/packages/backend/server/src/prelude.ts @@ -1,17 +1,5 @@ -import { homedir } from 'node:os'; -import { join } from 'node:path'; - -import { config } from 'dotenv'; - -import { SERVER_FLAVOR } from './config/default'; - -if (SERVER_FLAVOR === 'selfhosted') { - config({ path: join(homedir(), '.affine', '.env') }); -} else { - config(); -} - import 'reflect-metadata'; +import './affine.env'; import './affine'; import './affine.config'; diff --git a/packages/backend/server/src/session.ts b/packages/backend/server/src/session.ts index e3ca61c4f6..878bf0e821 100644 --- a/packages/backend/server/src/session.ts +++ b/packages/backend/server/src/session.ts @@ -1,44 +1,13 @@ -import KeyvRedis from '@keyv/redis'; -import { - FactoryProvider, - Global, - Inject, - Injectable, - Module, -} from '@nestjs/common'; -import Redis from 'ioredis'; -import Keyv from 'keyv'; +import { Global, Injectable, Module } from '@nestjs/common'; -import { Config } from './config'; - -export const KeyvProvide = Symbol('KeyvProvide'); - -export const KeyvProvider: FactoryProvider = { - provide: KeyvProvide, - useFactory(config: Config) { - if (config.redis.enabled) { - return new Keyv({ - store: new KeyvRedis( - new Redis(config.redis.port, config.redis.host, { - username: config.redis.username, - password: config.redis.password, - db: config.redis.database + 2, - }) - ), - }); - } else { - return new Keyv(); - } - }, - inject: [Config], -}; +import { SessionCache } from './cache/instances'; @Injectable() export class SessionService { private readonly prefix = 'session:'; private readonly sessionTtl = 30 * 60 * 1000; // 30 min - constructor(@Inject(KeyvProvide) private readonly cache: Keyv) {} + constructor(private readonly cache: SessionCache) {} /** * get session @@ -46,7 +15,7 @@ export class SessionService { * @returns */ async get(key: string) { - return this.cache.get(this.prefix + key); + return this.cache.get(this.prefix + key); } /** @@ -57,7 +26,9 @@ export class SessionService { * @returns return true if success */ async set(key: string, value?: any, sessionTtl = this.sessionTtl) { - return this.cache.set(this.prefix + key, value, sessionTtl); + return this.cache.set(this.prefix + key, value, { + ttl: sessionTtl, + }); } async delete(key: string) { @@ -67,7 +38,7 @@ export class SessionService { @Global() @Module({ - providers: [KeyvProvider, SessionService], - exports: [KeyvProvider, SessionService], + providers: [SessionService], + exports: [SessionService], }) export class SessionModule {} diff --git a/packages/backend/server/src/throttler.ts b/packages/backend/server/src/throttler.ts index 76cf0e5720..a94fe9355b 100644 --- a/packages/backend/server/src/throttler.ts +++ b/packages/backend/server/src/throttler.ts @@ -5,42 +5,50 @@ import { ThrottlerGuard, ThrottlerModule, ThrottlerModuleOptions, + ThrottlerOptionsFactory, } from '@nestjs/throttler'; -import Redis from 'ioredis'; import { ThrottlerStorageRedisService } from 'nestjs-throttler-storage-redis'; +import { ThrottlerCache } from './cache'; import { Config } from './config'; import { getRequestResponseFromContext } from './utils/nestjs'; +@Injectable() +class CustomOptionsFactory implements ThrottlerOptionsFactory { + constructor( + private readonly config: Config, + private readonly cache: ThrottlerCache + ) {} + createThrottlerOptions() { + const options: ThrottlerModuleOptions = { + throttlers: [ + { + ttl: this.config.rateLimiter.ttl, + limit: this.config.rateLimiter.limit, + }, + ], + skipIf: () => { + return !this.config.node.prod || this.config.affine.canary; + }, + }; + + if (this.config.redis.enabled) { + new Logger(RateLimiterModule.name).log('Use Redis'); + options.storage = new ThrottlerStorageRedisService( + // @ts-expect-error hidden field + this.cache.redis + ); + } + + return options; + } +} + @Global() @Module({ imports: [ ThrottlerModule.forRootAsync({ - inject: [Config], - useFactory: (config: Config): ThrottlerModuleOptions => { - const options: ThrottlerModuleOptions = { - throttlers: [ - { - ttl: config.rateLimiter.ttl, - limit: config.rateLimiter.limit, - }, - ], - skipIf: () => { - return !config.node.prod || config.affine.canary; - }, - }; - if (config.redis.enabled) { - new Logger(RateLimiterModule.name).log('Use Redis'); - options.storage = new ThrottlerStorageRedisService( - new Redis(config.redis.port, config.redis.host, { - username: config.redis.username, - password: config.redis.password, - db: config.redis.database + 1, - }) - ); - } - return options; - }, + useClass: CustomOptionsFactory, }), ], }) diff --git a/packages/backend/server/tests/auth.spec.ts b/packages/backend/server/tests/auth.spec.ts index b3edd260f2..e87125b796 100644 --- a/packages/backend/server/tests/auth.spec.ts +++ b/packages/backend/server/tests/auth.spec.ts @@ -3,6 +3,7 @@ import { Test, TestingModule } from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import test from 'ava'; +import { CacheModule } from '../src/cache'; import { ConfigModule } from '../src/config'; import { RevertCommand, RunCommand } from '../src/data/commands/run'; import { GqlModule } from '../src/graphql.module'; @@ -39,6 +40,7 @@ test.beforeEach(async () => { https: true, }), PrismaModule, + CacheModule, GqlModule, AuthModule, RateLimiterModule, diff --git a/packages/backend/server/tests/feature.spec.ts b/packages/backend/server/tests/feature.spec.ts index 50896d21d1..320f970898 100644 --- a/packages/backend/server/tests/feature.spec.ts +++ b/packages/backend/server/tests/feature.spec.ts @@ -5,6 +5,7 @@ import { Test, TestingModule } from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import ava, { type TestFn } from 'ava'; +import { CacheModule } from '../src/cache'; import { ConfigModule } from '../src/config'; import { RevertCommand, RunCommand } from '../src/data/commands/run'; import { AuthModule } from '../src/modules/auth'; @@ -80,6 +81,7 @@ test.beforeEach(async t => { }, }), PrismaModule, + CacheModule, AuthModule, FeatureModule, RateLimiterModule, diff --git a/packages/backend/server/tests/mailer.e2e.ts b/packages/backend/server/tests/mailer.e2e.ts index ea8beea45f..3538d4e303 100644 --- a/packages/backend/server/tests/mailer.e2e.ts +++ b/packages/backend/server/tests/mailer.e2e.ts @@ -10,6 +10,7 @@ import { Test, TestingModule } from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import ava, { type TestFn } from 'ava'; +import { CacheModule } from '../src/cache'; import { ConfigModule } from '../src/config'; import { RevertCommand, RunCommand } from '../src/data/commands/run'; import { GqlModule } from '../src/graphql.module'; @@ -45,6 +46,7 @@ test.beforeEach(async t => { PrismaModule, GqlModule, AuthModule, + CacheModule, RateLimiterModule, ], providers: [RevertCommand, RunCommand], diff --git a/packages/backend/server/tests/quota.spec.ts b/packages/backend/server/tests/quota.spec.ts index d819c21752..5e5237759a 100644 --- a/packages/backend/server/tests/quota.spec.ts +++ b/packages/backend/server/tests/quota.spec.ts @@ -4,6 +4,7 @@ import { Test, TestingModule } from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import ava, { type TestFn } from 'ava'; +import { CacheModule } from '../src/cache'; import { ConfigModule } from '../src/config'; import { RevertCommand, RunCommand } from '../src/data/commands/run'; import { EventModule } from '../src/event'; @@ -49,6 +50,7 @@ test.beforeEach(async t => { https: true, }), PrismaModule, + CacheModule, AuthModule, EventModule, QuotaModule, diff --git a/packages/backend/server/tests/session.spec.ts b/packages/backend/server/tests/session.spec.ts index 18d680e87d..71b08a485a 100644 --- a/packages/backend/server/tests/session.spec.ts +++ b/packages/backend/server/tests/session.spec.ts @@ -3,6 +3,7 @@ import { Test, TestingModule } from '@nestjs/testing'; import ava, { type TestFn } from 'ava'; +import { CacheModule } from '../src/cache'; import { ConfigModule } from '../src/config'; import { SessionModule, SessionService } from '../src/session'; @@ -19,6 +20,7 @@ test.beforeEach(async t => { enabled: false, }, }), + CacheModule, SessionModule, ], }).compile();