feat: auth metric and trace (#4063)

apps/core/src/utils/cloud-utils.tsx

@@ -1,11 +1,53 @@
 import { isDesktop } from '@affine/env/constant';
+import {
+  generateRandUTF16Chars,
+  SPAN_ID_BYTES,
+  TRACE_ID_BYTES,
+  traceReporter,
+} from '@affine/graphql';
 import { refreshRootMetadataAtom } from '@affine/workspace/atom';
 import { getCurrentStore } from '@toeverything/infra/atom';
 // eslint-disable-next-line @typescript-eslint/no-restricted-imports
 import { signIn, signOut } from 'next-auth/react';
 import { startTransition } from 'react';
 
+type TraceParams = {
+  startTime: string;
+  spanId: string;
+  traceId: string;
+  event: string;
+};
+
+function genTraceParams(): TraceParams {
+  const startTime = new Date().toISOString();
+  const spanId = generateRandUTF16Chars(SPAN_ID_BYTES);
+  const traceId = generateRandUTF16Chars(TRACE_ID_BYTES);
+  const event = 'signInCloud';
+  return { startTime, spanId, traceId, event };
+}
+
+function onResolveHandleTrace<T>(
+  res: Promise<T> | T,
+  params: TraceParams
+): Promise<T> | T {
+  const { startTime, spanId, traceId, event } = params;
+  traceReporter &&
+    traceReporter.cacheTrace(traceId, spanId, startTime, { event });
+  return res;
+}
+
+function onRejectHandleTrace<T>(
+  res: Promise<T> | T,
+  params: TraceParams
+): Promise<T> {
+  const { startTime, spanId, traceId, event } = params;
+  traceReporter &&
+    traceReporter.uploadTrace(traceId, spanId, startTime, { event });
+  return Promise.reject(res);
+}
+
 export const signInCloud: typeof signIn = async (provider, ...rest) => {
+  const traceParams = genTraceParams();
   if (isDesktop) {
     if (provider === 'google') {
       open(
@@ -29,25 +71,32 @@ export const signInCloud: typeof signIn = async (provider, ...rest) => {
           callbackUrl: buildCallbackUrl(callbackUrl),
         },
         ...tail
-      );
+      )
+        .then(res => onResolveHandleTrace(res, traceParams))
+        .catch(err => onRejectHandleTrace(err, traceParams));
     }
   } else {
-    return signIn(provider, ...rest);
+    return signIn(provider, ...rest)
+      .then(res => onResolveHandleTrace(res, traceParams))
+      .catch(err => onRejectHandleTrace(err, traceParams));
   }
 };
 
 export const signOutCloud: typeof signOut = async options => {
+  const traceParams = genTraceParams();
   return signOut({
     ...options,
     callbackUrl: '/',
-  }).then(result => {
-    if (result) {
-      startTransition(() => {
-        getCurrentStore().set(refreshRootMetadataAtom);
-      });
-    }
-    return result;
-  });
+  })
+    .then(result => {
+      if (result) {
+        startTransition(() => {
+          getCurrentStore().set(refreshRootMetadataAtom);
+        });
+      }
+      return onResolveHandleTrace(result, traceParams);
+    })
+    .catch(err => onRejectHandleTrace(err, traceParams));
 };
 
 export function buildCallbackUrl(callbackUrl: string) {

apps/server/src/metrics/metrics.ts

@@ -22,4 +22,7 @@ export class Metrics implements OnModuleDestroy {
   jwstCodecMerge = metricsCreator.counter('jwst_codec_merge');
   jwstCodecDidnotMatch = metricsCreator.counter('jwst_codec_didnot_match');
   jwstCodecFail = metricsCreator.counter('jwst_codec_fail');
+
+  authCounter = metricsCreator.counter('auth');
+  authFailCounter = metricsCreator.counter('auth_fail', ['reason']);
 }

apps/server/src/modules/auth/next-auth.controller.ts

@@ -19,6 +19,7 @@ import type { AuthAction, NextAuthOptions } from 'next-auth';
 import { AuthHandler } from 'next-auth/core';
 
 import { Config } from '../../config';
+import { Metrics } from '../../metrics/metrics';
 import { PrismaService } from '../../prisma/service';
 import { CloudThrottlerGuard, Throttle } from '../../throttler';
 import { NextAuthOptionsProvide } from './next-auth-options';
@@ -37,7 +38,8 @@ export class NextAuthController {
     readonly prisma: PrismaService,
     private readonly authService: AuthService,
     @Inject(NextAuthOptionsProvide)
-    private readonly nextAuthOptions: NextAuthOptions
+    private readonly nextAuthOptions: NextAuthOptions,
+    private readonly metrics: Metrics
   ) {
     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
     this.callbackSession = nextAuthOptions.callbacks!.session;
@@ -52,6 +54,7 @@ export class NextAuthController {
     @Query() query: Record<string, any>,
     @Next() next: NextFunction
   ) {
+    this.metrics.authCounter(1, {});
     const [action, providerId] = req.url // start with request url
       .slice(BASE_URL.length) // make relative to baseUrl
       .replace(/\?.*/, '') // remove query part, use only path part
@@ -83,6 +86,7 @@ export class NextAuthController {
     const options = this.nextAuthOptions;
     if (req.method === 'POST' && action === 'session') {
       if (typeof req.body !== 'object' || typeof req.body.data !== 'object') {
+        this.metrics.authFailCounter(1, { reason: 'invalid_session_data' });
         throw new BadRequestException(`Invalid new session data`);
       }
       const user = await this.updateSession(req, req.body.data);
@@ -130,6 +134,9 @@ export class NextAuthController {
       if (!req.headers?.referer) {
         res.redirect('https://community.affine.pro/c/insider-general/');
       } else {
+        this.metrics.authFailCounter(1, {
+          reason: 'no_early_access_permission',
+        });
         res.status(403);
         res.json({
           url: 'https://community.affine.pro/c/insider-general/',