feat: auth metric and trace (#4063)

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}`
   );