@finos/legend-extension-tracer-zipkin/src/ZipkinTracerPlugin.ts

Legend extension for tracer using Zipkin and OpenTracing

/**
 * Copyright (c) 2020-present, Goldman Sachs
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import packageJson from '../package.json' with { type: 'json' };
import { BatchRecorder, jsonEncoder } from 'zipkin';
import { HttpLogger } from 'zipkin-transport-http';
import type { Span as ZipkinSpan } from 'opentracing';
import {
  type TraceData,
  CORE_TRACER_TAG,
  assertNonEmptyString,
  guaranteeNonNullable,
  isNonNullable,
  TracerServicePlugin,
  HttpHeader,
  ContentType,
  CHARSET,
  type PlainObject,
} from '@finos/legend-shared';

/**
 * Previously, these exports rely on ES module interop to expose `default` export
 * properly. But since we use `ESM` for Typescript resolution now, we lose this
 * so we have to workaround by importing these and re-export them from CJS
 *
 * TODO: remove these when the package properly work with Typescript's nodenext
 * module resolution
 *
 * @workaround ESM
 * See https://github.com/microsoft/TypeScript/issues/49298
 */
import { default as SpanBuilder } from './CJS__Zipkin.cjs';
import type { default as ZipkinSpanBuilder } from 'zipkin-javascript-opentracing';

type ZipkinTracerPluginConfigData = {
  url: string;
  serviceName: string;
};

export class ZipkinTracerPlugin extends TracerServicePlugin<ZipkinSpan> {
  private _spanBuilder?: ZipkinSpanBuilder;

  constructor() {
    super(packageJson.extensions.tracerPlugin, packageJson.version);
  }

  override configure(
    configData: ZipkinTracerPluginConfigData,
  ): TracerServicePlugin<ZipkinSpan> {
    assertNonEmptyString(
      configData.url,
      `Can't configure Zipkin tracer: 'url' field is missing or empty`,
    );
    assertNonEmptyString(
      configData.serviceName,
      `Can't configure Zipkin tracer: 'serviceName' field is missing or empty`,
    );
    this._spanBuilder = new SpanBuilder.Zipkin({
      recorder: new BatchRecorder({
        logger: new HttpLogger({
          endpoint: configData.url,
          jsonEncoder: jsonEncoder.JSON_V2,
          // NOTE: this fetch implementation will be used for sending `spans`.
          // with some specific options, we have to customize this instead of using the default global fetch
          // See https://github.com/openzipkin/zipkin-js/tree/master/packages/zipkin-transport-http#optional
          fetchImplementation: (_url: string, options: PlainObject) =>
            fetch(_url, {
              ...options,
              mode: 'cors', // allow CORS - See https://developer.mozilla.org/en-US/docs/Web/API/Request/mode
              credentials: 'include', // allow sending credentials to other domain
              redirect: 'manual', // avoid following authentication redirects
              headers: {
                [HttpHeader.CONTENT_TYPE]: `${ContentType.APPLICATION_JSON};${CHARSET}`,
                [HttpHeader.ACCEPT]: ContentType.APPLICATION_JSON,
              },
            }),
          // eslint-disable-next-line @typescript-eslint/no-explicit-any
        } as any),
      }),
      serviceName: configData.serviceName,
      kind: 'client',
    });
    return this;
  }

  get spanBuilder(): ZipkinSpanBuilder {
    return guaranteeNonNullable(
      this._spanBuilder,
      `Can't configure Zipkin tracer: Tracer service has not been configured`,
    );
  }

  bootstrap(clientSpan: ZipkinSpan | undefined, response: Response): void {
    clientSpan?.setTag(
      CORE_TRACER_TAG.HTTP_STATUS,
      `${response.status} (${response.statusText})`,
    );
  }

  createClientSpan(
    traceData: TraceData,
    method: string,
    url: string,
    headers: PlainObject = {},
  ): ZipkinSpan {
    // When the service (client) calls a downstream services (server), it’s useful to pass down the SpanContext, so that
    // Spans generated by this service could join the Spans from our service in a single trace. To do that, our service needs to
    // `inject` the SpanContext into the payload and the downstream services need to `extract` the context info create more Spans
    // See https://opentracing.io/guides/java/inject-extract/
    // See https://github.com/DanielMSchmidt/zipkin-javascript-opentracing
    const clientSpan = this.spanBuilder.startSpan(traceData.name) as ZipkinSpan;
    if (traceData.tags) {
      Object.entries(traceData.tags).forEach(([tag, value]) => {
        if (isNonNullable(value)) {
          clientSpan.setTag(tag, value);
        }
      });
    }
    clientSpan.setTag(CORE_TRACER_TAG.HTTP_REQUEST_METHOD, method);
    clientSpan.setTag(CORE_TRACER_TAG.HTTP_REQUEST_URL, url);
    this.spanBuilder.inject(
      clientSpan,
      SpanBuilder.Zipkin.FORMAT_HTTP_HEADERS,
      headers,
    );
    return clientSpan;
  }

  concludeClientSpan(
    clientSpan: ZipkinSpan | undefined,
    error: Error | undefined,
  ): void {
    if (!clientSpan) {
      return;
    }
    if (error) {
      clientSpan.setTag(CORE_TRACER_TAG.RESULT, 'error');
      if (error.message) {
        clientSpan.setTag(CORE_TRACER_TAG.ERROR, error.message);
      }
    } else {
      clientSpan.setTag(CORE_TRACER_TAG.RESULT, 'success');
    }
    clientSpan.finish();
  }
}