@cerbos/opentelemetry/src/transport.ts

OpenTelemetry instrumentation for the @cerbos/grpc and @cerbos/http client libraries

import type {
  DescMessage,
  DescMethod,
  DescMethodServerStreaming,
  DescMethodUnary,
  MessageShape,
  MessageValidType,
} from "@bufbuild/protobuf";
import type { Attributes, SpanStatus, Tracer } from "@opentelemetry/api";
import {
  SpanKind,
  SpanStatusCode,
  context,
  propagation,
  trace,
} from "@opentelemetry/api";
import {
  ATTR_RPC_GRPC_STATUS_CODE,
  ATTR_RPC_METHOD,
  ATTR_RPC_SERVICE,
  ATTR_RPC_SYSTEM,
} from "@opentelemetry/semantic-conventions/incubating";

import { NotOK, Status } from "@cerbos/core";
import type {
  AbortHandler,
  Transport as CoreTransport,
} from "@cerbos/core/~internal";
import { methodName } from "@cerbos/core/~internal";

import type { CerbosInstrumentation } from "./instrumentation.js";
import type { Instruments } from "./instruments.js";

export class Transport implements CoreTransport {
  private readonly transport: {
    [MethodKind in keyof CoreTransport]: CoreTransport[MethodKind];
  };

  public constructor(
    private readonly instrumentation: CerbosInstrumentation,
    transport: CoreTransport,
  ) {
    this.transport = {
      unary: transport.unary.bind(transport),
      serverStream: transport.serverStream.bind(transport),
    };
  }

  public async unary<I extends DescMessage, O extends DescMessage>(
    method: DescMethodUnary<I, O>,
    request: MessageValidType<I>,
    headers: Headers,
    abortHandler: AbortHandler,
  ): Promise<MessageShape<O>> {
    const { call, succeeded, failed } = this.instrument(
      this.transport.unary,
      method,
      headers,
    );

    try {
      const output = await call(method, request, headers, abortHandler);
      succeeded();
      return output;
    } catch (error) {
      failed(error);
      throw error;
    }
  }

  public async *serverStream<I extends DescMessage, O extends DescMessage>(
    method: DescMethodServerStreaming<I, O>,
    request: MessageValidType<I>,
    headers: Headers,
    abortHandler: AbortHandler,
  ): AsyncGenerator<MessageShape<O>, void, undefined> {
    const { call, succeeded, failed } = this.instrument(
      this.transport.serverStream,
      method,
      headers,
    );

    let done = false;

    try {
      yield* call(method, request, headers, abortHandler);
      done = true;
      succeeded();
    } catch (error) {
      done = true;
      failed(error);
      throw error;
    } finally {
      if (!done) {
        failed(abortHandler.error());
      }
    }
  }

  private instrument<F>(
    fn: F,
    method: DescMethod,
    headers: Headers,
  ): {
    call: F;
    succeeded: () => void;
    failed: (error: unknown) => void;
  } {
    const startTime = performance.now();

    const status: SpanStatus = { code: SpanStatusCode.UNSET };

    const attributes: Attributes = {
      [ATTR_RPC_SYSTEM]: "grpc",
      [ATTR_RPC_SERVICE]: method.parent.typeName,
      [ATTR_RPC_METHOD]: method.name,
    };

    const span = this.tracer.startSpan(methodName(method), {
      kind: SpanKind.CLIENT,
      startTime,
    });

    const activeContext = trace.setSpan(context.active(), span);
    propagation.inject(activeContext, headers, {
      set(carrier, key, value) {
        carrier.set(key, value);
      },
    });

    const finish = (): void => {
      const endTime = performance.now();
      span.setStatus(status);
      span.setAttributes(attributes);
      span.end(endTime);
      this.instruments.duration.record(endTime - startTime, attributes);
    };

    return {
      call: context.bind(activeContext, fn),
      succeeded: (): void => {
        attributes[ATTR_RPC_GRPC_STATUS_CODE] = Status.OK;
        finish();
      },
      failed: (error: unknown): void => {
        status.code = SpanStatusCode.ERROR;

        if (error instanceof Error) {
          status.message = error.message;
          attributes["cerbos.error"] = error.message;

          if (error instanceof NotOK) {
            attributes[ATTR_RPC_GRPC_STATUS_CODE] = error.code;
          }
        }

        finish();
      },
    };
  }

  private get instruments(): Instruments {
    return this.instrumentation["~instruments"];
  }

  private get tracer(): Tracer {
    return this.instrumentation["~tracer"];
  }
}