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/',

packages/graphql/src/__tests__/fetcher.spec.ts

@@ -143,8 +143,8 @@ describe('Trace Reporter', () => {
     const traceSpan = TraceReporter.createTraceSpan(
       traceId,
       spanId,
-      requestId,
-      startTime
+      startTime,
+      { requestId }
     );
     expect(traceSpan.startTime).toBe(startTime);
     expect(
@@ -152,7 +152,7 @@ describe('Trace Reporter', () => {
         `projects/{GCP_PROJECT_ID}/traces/${traceId}/spans/${spanId}`
     ).toBe(true);
     expect(traceSpan.spanId).toBe(spanId);
-    expect(traceSpan.attributes.attributeMap.requestId.stringValue.value).toBe(
+    expect(traceSpan.attributes.attributeMap.requestId?.stringValue.value).toBe(
       requestId
     );
   });

packages/graphql/src/fetcher.ts

@@ -181,13 +181,14 @@ export const gqlFetcherFactory = (endpoint: string) => {
     if (!isFormData) {
       headers['content-type'] = 'application/json';
     }
-    const ret = fetchWithReport(
+    const ret = fetchWithTraceReport(
       endpoint,
       merge(options.context, {
         method: 'POST',
         headers,
         body: isFormData ? body : JSON.stringify(body),
-      })
+      }),
+      { event: 'GraphQLRequest' }
     ).then(async res => {
       if (res.headers.get('content-type')?.startsWith('application/json')) {
         const result = (await res.json()) as ExecutionResult;
@@ -214,9 +215,10 @@ export const gqlFetcherFactory = (endpoint: string) => {
   return gqlFetch;
 };
 
-export const fetchWithReport = (
+export const fetchWithTraceReport = (
   input: RequestInfo | URL,
-  init?: RequestInit
+  init?: RequestInit,
+  traceOptions?: { event: string }
 ): Promise<Response> => {
   const startTime = new Date().toISOString();
   const spanId = generateRandUTF16Chars(SPAN_ID_BYTES);
@@ -225,6 +227,7 @@ export const fetchWithReport = (
   init = init || {};
   init.headers = init.headers || new Headers();
   const requestId = nanoid();
+  const event = traceOptions?.event;
   if (init.headers instanceof Headers) {
     init.headers.append('x-request-id', requestId);
     init.headers.append('traceparent', traceparent);
@@ -241,12 +244,18 @@ export const fetchWithReport = (
   return fetch(input, init)
     .then(response => {
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      traceReporter!.cacheTrace(traceId, spanId, requestId, startTime);
+      traceReporter!.cacheTrace(traceId, spanId, startTime, {
+        requestId,
+        ...(event ? { event } : {}),
+      });
       return response;
     })
     .catch(err => {
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      traceReporter!.uploadTrace(traceId, spanId, requestId, startTime);
+      traceReporter!.uploadTrace(traceId, spanId, startTime, {
+        requestId,
+        ...(event ? { event } : {}),
+      });
       return Promise.reject(err);
     });
 };

packages/graphql/src/index.ts

@@ -1,4 +1,5 @@
 export * from './fetcher';
 export * from './graphql';
 export * from './schema';
+export * from './utils';
 import '@affine/env/global';

packages/graphql/src/utils.ts

@@ -16,12 +16,18 @@ type TraceSpan = {
   endTime: string;
   attributes: {
     attributeMap: {
-      requestId: {
+      requestId?: {
         stringValue: {
           value: string;
           truncatedByteCount: number;
         };
       };
+      event?: {
+        stringValue: {
+          value: string;
+          truncatedByteCount: 0;
+        };
+      };
     };
     droppedAttributesCount: number;
   };
@@ -65,14 +71,17 @@ export class TraceReporter {
   public cacheTrace(
     traceId: string,
     spanId: string,
-    requestId: string,
-    startTime: string
+    startTime: string,
+    attributes: {
+      requestId?: string;
+      event?: string;
+    }
   ) {
     const span = TraceReporter.createTraceSpan(
       traceId,
       spanId,
-      requestId,
-      startTime
+      startTime,
+      attributes
     );
     this.spansCache.push(span);
     if (this.spansCache.length <= 1) {
@@ -83,14 +92,17 @@ export class TraceReporter {
   public uploadTrace(
     traceId: string,
     spanId: string,
-    requestId: string,
-    startTime: string
+    startTime: string,
+    attributes: {
+      requestId?: string;
+      event?: string;
+    }
   ) {
     const span = TraceReporter.createTraceSpan(
       traceId,
       spanId,
-      requestId,
-      startTime
+      startTime,
+      attributes
     );
     TraceReporter.reportToTraceEndpoint(JSON.stringify({ spans: [span] }));
   }
@@ -114,26 +126,46 @@ export class TraceReporter {
   public static createTraceSpan(
     traceId: string,
     spanId: string,
-    requestId: string,
-    startTime: string
+    startTime: string,
+    attributes: {
+      requestId?: string;
+      event?: string;
+    }
   ): TraceSpan {
+    const requestId = attributes.requestId;
+    const event = attributes.event;
+
     return {
       name: `projects/{GCP_PROJECT_ID}/traces/${traceId}/spans/${spanId}`,
       spanId,
       displayName: {
-        value: 'fetch',
+        value: 'AFFiNE_REQUEST',
         truncatedByteCount: 0,
       },
       startTime,
       endTime: new Date().toISOString(),
       attributes: {
         attributeMap: {
-          requestId: {
-            stringValue: {
-              value: requestId,
-              truncatedByteCount: 0,
-            },
-          },
+          ...(!requestId
+            ? {}
+            : {
+                requestId: {
+                  stringValue: {
+                    value: requestId,
+                    truncatedByteCount: 0,
+                  },
+                },
+              }),
+          ...(!event
+            ? {}
+            : {
+                event: {
+                  stringValue: {
+                    value: event,
+                    truncatedByteCount: 0,
+                  },
+                },
+              }),
         },
         droppedAttributesCount: 0,
       },

packages/workspace/src/blob/cloud-blob-storage.ts

@@ -1,6 +1,6 @@
 import {
   deleteBlobMutation,
-  fetchWithReport,
+  fetchWithTraceReport,
   listBlobsQuery,
   setBlobMutation,
 } from '@affine/graphql';
@@ -12,7 +12,7 @@ export const createCloudBlobStorage = (workspaceId: string): BlobStorage => {
   return {
     crud: {
       get: async key => {
-        return fetchWithReport(
+        return fetchWithTraceReport(
           runtimeConfig.serverUrlPrefix +
             `/api/workspaces/${workspaceId}/blobs/${key}`
         ).then(res => res.blob());

packages/workspace/src/providers/cloud/index.ts

@@ -1,5 +1,5 @@
 import { DebugLogger } from '@affine/debug';
-import { fetchWithReport } from '@affine/graphql';
+import { fetchWithTraceReport } from '@affine/graphql';
 import type { ActiveDocProvider, DocProviderCreator } from '@blocksuite/store';
 import { Workspace } from '@blocksuite/store';
 import type { Doc } from 'yjs';
@@ -17,7 +17,7 @@ export async function downloadBinaryFromCloud(
   if (hashMap.has(`${rootGuid}/${pageGuid}`)) {
     return true;
   }
-  const response = await fetchWithReport(
+  const response = await fetchWithTraceReport(
     runtimeConfig.serverUrlPrefix +
       `/api/workspaces/${rootGuid}/docs/${pageGuid}`
   );