fetch.ts

import { AsyncLocalStorage } from 'node:async_hooks';
import { debuglog } from 'node:util';

import { fetch as UndiciFetch, Request, Response, Agent, getGlobalDispatcher, Pool, Dispatcher } from 'undici';
import type { RequestInfo, RequestInit } from 'undici';
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
import undiciSymbols from 'undici/lib/core/symbols.js';
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
import { getResponseState } from 'undici/lib/web/fetch/response.js';

import { BaseAgent } from './BaseAgent.js';
import type { BaseAgentOptions } from './BaseAgent.js';
import { initDiagnosticsChannel } from './diagnosticsChannel.js';
import type { FetchOpaque } from './FetchOpaqueInterceptor.js';
import { HttpAgent } from './HttpAgent.js';
import type { HttpAgentOptions } from './HttpAgent.js';
import { channels } from './HttpClient.js';
import type {
  ClientOptions,
  PoolStat,
  RequestDiagnosticsMessage,
  ResponseDiagnosticsMessage,
  UndiciTimingInfo,
} from './HttpClient.js';
import type { IncomingHttpHeaders } from './IncomingHttpHeaders.js';
import type { FetchMeta, HttpMethod, RequestMeta } from './Request.js';
import type { RawResponseWithMeta, SocketInfo } from './Response.js';
import symbols from './symbols.js';
import { convertHeader, globalId, performanceTime, updateSocketInfo } from './utils.js';

const debug = debuglog('urllib/fetch');

export interface UrllibRequestInit extends RequestInit {
  // default is true
  timing?: boolean;
}

export type FetchDiagnosticsMessage = {
  fetch: FetchMeta;
  fetchOpaque: FetchOpaque;
};

export type FetchResponseDiagnosticsMessage = {
  fetch: FetchMeta;
  fetchOpaque: FetchOpaque;
  timingInfo?: UndiciTimingInfo;
  response?: Response;
  error?: Error;
};

export class FetchFactory {
  #dispatcher?: Dispatcher.ComposedDispatcher;
  #opaqueLocalStorage = new AsyncLocalStorage<FetchOpaque>();

  static #instance = new FetchFactory();

  setClientOptions(clientOptions: ClientOptions): void {
    let dispatcherOption: BaseAgentOptions = {
      opaqueLocalStorage: this.#opaqueLocalStorage,
    };
    let dispatcherClazz: new (options: BaseAgentOptions) => BaseAgent = BaseAgent;
    if (clientOptions?.lookup || clientOptions?.checkAddress) {
      dispatcherOption = {
        ...dispatcherOption,
        lookup: clientOptions.lookup,
        checkAddress: clientOptions.checkAddress,
        connect: clientOptions.connect,
        allowH2: clientOptions.allowH2,
      } as HttpAgentOptions;
      dispatcherClazz = HttpAgent as unknown as new (options: BaseAgentOptions) => BaseAgent;
    } else if (clientOptions?.connect) {
      dispatcherOption = {
        ...dispatcherOption,
        connect: clientOptions.connect,
        allowH2: clientOptions.allowH2,
      } as HttpAgentOptions;
      dispatcherClazz = BaseAgent;
    } else if (clientOptions?.allowH2) {
      // Support HTTP2
      dispatcherOption = {
        ...dispatcherOption,
        allowH2: clientOptions.allowH2,
      } as HttpAgentOptions;
      dispatcherClazz = BaseAgent;
    }
    this.#dispatcher = new dispatcherClazz(dispatcherOption);
    initDiagnosticsChannel();
  }

  getDispatcher(): Dispatcher {
    return this.#dispatcher ?? getGlobalDispatcher();
  }

  setDispatcher(dispatcher: Agent): void {
    this.#dispatcher = dispatcher;
  }

  getDispatcherPoolStats(): Record<string, PoolStat> {
    const agent = this.getDispatcher();
    // origin => Pool Instance
    const clients: Map<string, WeakRef<Pool>> | undefined = Reflect.get(agent, undiciSymbols.kClients);
    const poolStatsMap: Record<string, PoolStat> = {};
    if (!clients) {
      return poolStatsMap;
    }
    for (const [key, ref] of clients) {
      const pool = (typeof ref.deref === 'function' ? ref.deref() : ref) as unknown as Pool & { dispatcher: Pool };
      // NOTE: pool become to { dispatcher: Pool } in undici@v7
      const stats = pool?.stats ?? pool?.dispatcher?.stats;
      if (!stats) continue;

      poolStatsMap[key] = {
        connected: stats.connected,
        free: stats.free,
        pending: stats.pending,
        queued: stats.queued,
        running: stats.running,
        size: stats.size,
      } satisfies PoolStat;
    }
    return poolStatsMap;
  }

  static setClientOptions(clientOptions: ClientOptions): void {
    FetchFactory.#instance.setClientOptions(clientOptions);
  }

  static getDispatcherPoolStats(): Record<string, PoolStat> {
    return FetchFactory.#instance.getDispatcherPoolStats();
  }

  async fetch(input: RequestInfo, init?: UrllibRequestInit): Promise<Response> {
    const requestStartTime = performance.now();
    init = init ?? {};
    init.dispatcher = init.dispatcher ?? this.#dispatcher;
    const request = new Request(input, init);
    const requestId = globalId('HttpClientRequest');
    // https://developer.chrome.com/docs/devtools/network/reference/?utm_source=devtools#timing-explanation
    const timing = {
      // socket assigned
      queuing: 0,
      // dns lookup time
      dnslookup: 0,
      // socket connected
      connected: 0,
      // request headers sent
      requestHeadersSent: 0,
      // request sent, including headers and body
      requestSent: 0,
      // Time to first byte (TTFB), the response headers have been received
      waiting: 0,
      // the response body and trailers have been received
      contentDownload: 0,
    };

    // using opaque to diagnostics channel, binding request and socket
    const internalOpaque = {
      [symbols.kRequestId]: requestId,
      [symbols.kRequestStartTime]: requestStartTime,
      [symbols.kEnableRequestTiming]: !!(init.timing ?? true),
      [symbols.kRequestTiming]: timing,
      // [symbols.kRequestOriginalOpaque]: originalOpaque,
    } as FetchOpaque;
    const reqMeta: RequestMeta = {
      requestId,
      url: request.url,
      args: {
        method: request.method as HttpMethod,
        type: request.method as HttpMethod,
        data: request.body,
        headers: convertHeader(request.headers),
      },
      retries: 0,
    };
    const fetchMeta: FetchMeta = {
      requestId,
      request,
    };
    const socketInfo: SocketInfo = {
      id: 0,
      localAddress: '',
      localPort: 0,
      remoteAddress: '',
      remotePort: 0,
      remoteFamily: '',
      bytesWritten: 0,
      bytesRead: 0,
      handledRequests: 0,
      handledResponses: 0,
    };
    channels.request.publish({
      request: reqMeta,
      isSentByFetch: true,
      fetchOpaque: internalOpaque,
    } as RequestDiagnosticsMessage);
    channels.fetchRequest.publish({
      fetch: fetchMeta,
      fetchOpaque: internalOpaque,
    } as FetchDiagnosticsMessage);

    let res: Response;
    // keep urllib createCallbackResponse style
    const resHeaders: IncomingHttpHeaders = {};
    const urllibResponse = {
      status: -1,
      statusCode: -1,
      statusText: '',
      statusMessage: '',
      headers: resHeaders,
      size: 0,
      aborted: false,
      rt: 0,
      keepAliveSocket: true,
      requestUrls: [request.url],
      timing,
      socket: socketInfo,
      retries: 0,
      socketErrorRetries: 0,
    } as any as RawResponseWithMeta;
    try {
      await this.#opaqueLocalStorage.run(internalOpaque, async () => {
        res = await UndiciFetch(request);
      });
    } catch (e: any) {
      updateSocketInfo(socketInfo, internalOpaque, e);
      urllibResponse.rt = performanceTime(requestStartTime);
      debug('Request#%d throw error: %s', requestId, e);
      channels.fetchResponse.publish({
        fetch: fetchMeta,
        error: e,
        fetchOpaque: internalOpaque,
      } as FetchResponseDiagnosticsMessage);
      channels.response.publish({
        request: reqMeta,
        response: urllibResponse,
        error: e,
        isSentByFetch: true,
        fetchOpaque: internalOpaque,
      } as ResponseDiagnosticsMessage);
      throw e;
    }

    // get undici internal response
    // Bun's Response doesn't have the same internal state as npm undici's Response
    let state: any;
    try {
      state = getResponseState(res!);
    } catch {
      state = {};
    }
    updateSocketInfo(socketInfo, internalOpaque);

    urllibResponse.headers = convertHeader(res!.headers);
    urllibResponse.status = urllibResponse.statusCode = res!.status;
    urllibResponse!.statusMessage = res!.statusText;
    if (urllibResponse.headers['content-length']) {
      urllibResponse.size = parseInt(urllibResponse.headers['content-length']);
    }
    urllibResponse.rt = performanceTime(requestStartTime);
    debug(
      'Request#%d got response, status: %s, headers: %j, timing: %j, socket: %j',
      requestId,
      urllibResponse.status,
      urllibResponse.headers,
      timing,
      urllibResponse.socket,
    );
    channels.fetchResponse.publish({
      fetch: fetchMeta,
      timingInfo: state.timingInfo,
      response: res!,
      fetchOpaque: internalOpaque,
    } as FetchResponseDiagnosticsMessage);
    channels.response.publish({
      request: reqMeta,
      response: urllibResponse,
      isSentByFetch: true,
      fetchOpaque: internalOpaque,
    } as ResponseDiagnosticsMessage);
    return res!;
  }

  static getDispatcher(): Dispatcher {
    return FetchFactory.#instance.getDispatcher();
  }

  static setDispatcher(dispatcher: Agent): void {
    FetchFactory.#instance.setDispatcher(dispatcher);
  }

  static async fetch(input: RequestInfo, init?: UrllibRequestInit): Promise<Response> {
    return FetchFactory.#instance.fetch(input, init);
  }
}

export const fetch: (input: RequestInfo, init?: UrllibRequestInit) => Promise<Response> = FetchFactory.fetch;