chore: initiate traceIds from client (#1210)

Author: juliusmarminge

packages/uploadthing/src/_internal/client-future.ts

@@ -17,6 +17,8 @@ import type {
   FileRouter as AnyFileRouter,
   NewPresignedUrl,
 } from "../types";
+import type { TraceHeaders } from "./random-hex";
+import { generateTraceHeaders } from "./random-hex";
 import type { UploadPutResult } from "./types";
 import { createUTReporter } from "./ut-reporter";
 
@@ -377,6 +379,7 @@ export interface UploadFileOptions<TRoute extends AnyFileRoute> {
   files: AnyFile<TRoute>[];
   input: TRoute["$types"]["input"];
   onEvent: (event: UploadEvent<TRoute>) => void;
+  traceHeaders: TraceHeaders;
 
   XHRImpl: new () => XMLHttpRequest;
 }
@@ -391,7 +394,7 @@ export function uploadFile<TRoute extends AnyFileRoute>(
   url: string,
   { file, files, XHRImpl, ...options }: UploadFileOptions<TRoute>,
 ): Micro.Micro<UploadedFile<TRoute> | FailedFile<TRoute>, never, FetchContext> {
-  return fetchEff(url, { method: "HEAD" }).pipe(
+  return fetchEff(url, { method: "HEAD", headers: options.traceHeaders }).pipe(
     Micro.map(({ headers }) =>
       Number.parseInt(headers.get("x-ut-range-start") ?? "0"),
     ),
@@ -411,6 +414,8 @@ export function uploadFile<TRoute extends AnyFileRoute>(
         const rangeStart = uploadingFile.sent;
         xhr.setRequestHeader("Range", `bytes=${rangeStart}-`);
         xhr.setRequestHeader("x-uploadthing-version", version);
+        xhr.setRequestHeader("b3", options.traceHeaders.b3);
+        xhr.setRequestHeader("traceparent", options.traceHeaders.traceparent);
         xhr.responseType = "json";
 
         xhr.upload.addEventListener("progress", (ev) => {
@@ -525,6 +530,10 @@ export interface RequestPresignedUrlsOptions<
    * @example { Authorization: "Bearer 123" }
    */
   headers?: HeadersInit | LazyArg<MaybePromise<HeadersInit>> | undefined;
+  /**
+   * Custom trace headers to send with the request
+   */
+  traceHeaders: TraceHeaders;
   /**
    * The uploadthing package that is making this request, used to identify the client in the server logs
    * @example "@uploadthing/react"
@@ -551,6 +560,7 @@ export function requestPresignedUrls<
     package: options.package,
     url: options.url,
     headers: options.headers,
+    traceHeaders: options.traceHeaders,
   });
 
   return reportEventToUT("upload", {
@@ -592,14 +602,15 @@ export function uploadFiles<
   TEndpoint extends keyof TRouter,
 >(endpoint: TEndpoint, options: UploadFilesOptions<TRouter[TEndpoint]>) {
   const pendingFiles = options.files.map(makePendingFile);
-
+  const traceHeaders = generateTraceHeaders();
   return requestPresignedUrls({
     endpoint: endpoint,
     files: options.files,
     url: options.url,
     input: options.input,
     headers: options.headers,
     package: options.package,
+    traceHeaders,
   }).pipe(
     Micro.map(Arr.zip(pendingFiles)),
     Micro.tap((pairs) => {
@@ -622,6 +633,7 @@ export function uploadFiles<
             input: options.input,
             onEvent: options.onEvent,
             XHRImpl: globalThis.XMLHttpRequest,
+            traceHeaders,
           }),
         { concurrency: 6 },
       ),

packages/uploadthing/src/_internal/random-hex.ts

@@ -0,0 +1,27 @@
+export const randomHexString = (function () {
+  const characters = "abcdef0123456789";
+  const charactersLength = characters.length;
+  return function (length: number) {
+    let result = "";
+    for (let i = 0; i < length; i++) {
+      result += characters.charAt(Math.floor(Math.random() * charactersLength));
+    }
+    return result;
+  };
+})();
+
+export type TraceHeaders = {
+  b3: string;
+  traceparent: string;
+};
+
+export const generateTraceHeaders = (): TraceHeaders => {
+  const traceId = randomHexString(32);
+  const spanId = randomHexString(16);
+  const sampled = "01";
+
+  return {
+    b3: `${traceId}-${spanId}-${sampled}`,
+    traceparent: `00-${traceId}-${spanId}-${sampled}`,
+  };
+};

packages/uploadthing/src/_internal/upload-browser.ts

@@ -14,28 +14,36 @@ import type {
   UploadFilesOptions,
 } from "../types";
 import { logDeprecationWarning } from "./deprecations";
+import type { TraceHeaders } from "./random-hex";
+import { generateTraceHeaders } from "./random-hex";
 import type { UploadPutResult } from "./types";
 import { createUTReporter } from "./ut-reporter";
 
 const uploadWithProgress = (
   file: File,
   rangeStart: number,
   presigned: NewPresignedUrl,
-  onUploadProgress?:
-    | ((opts: { loaded: number; delta: number }) => void)
-    | undefined,
+  opts: {
+    traceHeaders: TraceHeaders;
+    onUploadProgress?:
+      | ((opts: { loaded: number; delta: number }) => void)
+      | undefined;
+  },
 ) =>
   Micro.async<unknown, UploadThingError, FetchContext>((resume) => {
     const xhr = new XMLHttpRequest();
     xhr.open("PUT", presigned.url, true);
     xhr.setRequestHeader("Range", `bytes=${rangeStart}-`);
     xhr.setRequestHeader("x-uploadthing-version", version);
+    xhr.setRequestHeader("b3", opts.traceHeaders.b3);
+    xhr.setRequestHeader("traceparent", opts.traceHeaders.traceparent);
+
     xhr.responseType = "json";
 
     let previousLoaded = 0;
     xhr.upload.addEventListener("progress", ({ loaded }) => {
       const delta = loaded - previousLoaded;
-      onUploadProgress?.({ loaded, delta });
+      opts.onUploadProgress?.({ loaded, delta });
       previousLoaded = loaded;
     });
     xhr.addEventListener("load", () => {
@@ -108,13 +116,17 @@ export const uploadFile = <
   file: File,
   presigned: NewPresignedUrl,
   opts: {
+    traceHeaders: TraceHeaders;
     onUploadProgress?: (progressEvent: {
       loaded: number;
       delta: number;
     }) => void;
   },
 ) =>
-  fetchEff(presigned.url, { method: "HEAD" }).pipe(
+  fetchEff(presigned.url, {
+    method: "HEAD",
+    headers: opts.traceHeaders,
+  }).pipe(
     Micro.map(({ headers }) =>
       parseInt(headers.get("x-ut-range-start") ?? "0", 10),
     ),
@@ -125,12 +137,14 @@ export const uploadFile = <
       }),
     ),
     Micro.flatMap((start) =>
-      uploadWithProgress(file, start, presigned, (progressEvent) =>
-        opts.onUploadProgress?.({
-          delta: progressEvent.delta,
-          loaded: progressEvent.loaded + start,
-        }),
-      ),
+      uploadWithProgress(file, start, presigned, {
+        traceHeaders: opts.traceHeaders,
+        onUploadProgress: (progressEvent) =>
+          opts.onUploadProgress?.({
+            delta: progressEvent.delta,
+            loaded: progressEvent.loaded + start,
+          }),
+      }),
     ),
     Micro.map(unsafeCoerce<unknown, UploadPutResult<TServerOutput>>),
     Micro.map((uploadResponse) => ({
@@ -171,11 +185,13 @@ export const uploadFilesInternal = <
   FetchContext
 > => {
   // classic service right here
+  const traceHeaders = generateTraceHeaders();
   const reportEventToUT = createUTReporter({
     endpoint: String(endpoint),
     package: opts.package,
     url: opts.url,
     headers: opts.headers,
+    traceHeaders,
   });
 
   const totalSize = opts.files.reduce((acc, f) => acc + f.size, 0);
@@ -204,6 +220,7 @@ export const uploadFilesInternal = <
                 opts.files[i]!,
                 presigned,
                 {
+                  traceHeaders,
                   onUploadProgress: (ev) => {
                     totalLoaded += ev.delta;
                     opts.onUploadProgress?.({

packages/uploadthing/src/_internal/ut-reporter.ts

@@ -10,6 +10,7 @@ import {
 } from "@uploadthing/shared";
 
 import * as pkgJson from "../../package.json";
+import type { TraceHeaders } from "./random-hex";
 import type { ActionType } from "./shared-schemas";
 import type { UTEvents } from "./types";
 
@@ -48,6 +49,7 @@ export const createUTReporter =
     endpoint: string;
     package?: string | undefined;
     headers: HeadersInit | (() => MaybePromise<HeadersInit>) | undefined;
+    traceHeaders: TraceHeaders;
   }): UTReporter =>
   (type, payload) =>
     Micro.gen(function* () {
@@ -66,6 +68,8 @@ export const createUTReporter =
       }
       headers.set("x-uploadthing-version", pkgJson.version);
       headers.set("Content-Type", "application/json");
+      headers.set("b3", cfg.traceHeaders.b3);
+      headers.set("traceparent", cfg.traceHeaders.traceparent);
 
       const response = yield* fetchEff(url, {
         method: "POST",

packages/uploadthing/src/client-future.ts

@@ -27,6 +27,7 @@ import {
 } from "./_internal/client-future";
 import type { Deferred } from "./_internal/deferred";
 import { createDeferred } from "./_internal/deferred";
+import { generateTraceHeaders } from "./_internal/random-hex";
 import type {
   EndpointArg,
   FileRouter,
@@ -73,13 +74,15 @@ export const future_genUploader = <TRouter extends FileRouter>(
     const endpoint = typeof slug === "function" ? slug(routeRegistry) : slug;
     const fetchFn: FetchEsque = initOpts?.fetch ?? window.fetch;
 
+    const traceHeaders = generateTraceHeaders();
     const pExit = await requestPresignedUrls({
       endpoint: String(endpoint),
       files: options.files,
       url: resolveMaybeUrlArg(initOpts?.url),
       // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
       input: (options as any).input as inferEndpointInput<TRouter[TEndpoint]>,
       headers: options.headers,
+      traceHeaders,
     }).pipe(Micro.provideService(FetchContext, fetchFn), (effect) =>
       Micro.runPromiseExit(
         effect,
@@ -109,6 +112,7 @@ export const future_genUploader = <TRouter extends FileRouter>(
         files: pendingFiles,
         input: options.input,
         onEvent: options.onEvent,
+        traceHeaders,
         XHRImpl: globalThis.XMLHttpRequest,
       }).pipe(Micro.provideService(FetchContext, fetchFn));
 

packages/uploadthing/src/client.ts

@@ -16,6 +16,7 @@ import {
 import * as pkgJson from "../package.json";
 import type { Deferred } from "./_internal/deferred";
 import { createDeferred } from "./_internal/deferred";
+import { generateTraceHeaders } from "./_internal/random-hex";
 import { uploadFile, uploadFilesInternal } from "./_internal/upload-browser";
 import { createUTReporter } from "./_internal/ut-reporter";
 import type {
@@ -110,11 +111,13 @@ export const genUploader = <TRouter extends FileRouter>(
 
     const endpoint = typeof slug === "function" ? slug(routeRegistry) : slug;
 
+    const traceHeaders = generateTraceHeaders();
     const utReporter = createUTReporter({
       endpoint: String(endpoint),
       package: initOpts?.package ?? "uploadthing/client",
       url: resolveMaybeUrlArg(initOpts?.url),
       headers: opts.headers,
+      traceHeaders,
     });
     const fetchFn: FetchEsque = initOpts?.fetch ?? window.fetch;
 
@@ -136,6 +139,7 @@ export const genUploader = <TRouter extends FileRouter>(
 
     const uploadEffect = (file: File, presigned: NewPresignedUrl) =>
       uploadFile(file, presigned, {
+        traceHeaders,
         onUploadProgress: (progressEvent) => {
           totalLoaded += progressEvent.delta;
           opts.onUploadProgress?.({

packages/uploadthing/test/browser/client.test.ts

@@ -234,6 +234,8 @@ describe("uploadFiles", () => {
       Object.fromEntries(middlewareMock.mock.calls[0]![0]!.req.headers),
     ).toMatchObject({
       authorization: "Bearer my-auth-token",
+      b3: expect.any(String),
+      traceparent: expect.any(String),
       "x-uploadthing-package": "vitest",
       "x-uploadthing-version": expect.stringMatching(/\d+\.\d+\.\d+/),
     });