src/agent/protocol/grpc/clients/TraceReportClient.ts

/*!
 *
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 config from '../../../../config/AgentConfig';
import * as grpc from '@grpc/grpc-js';
import { connectivityState } from '@grpc/grpc-js';
import { createLogger } from '../../../../logging';
import Client from './Client';
import { TraceSegmentReportServiceClient } from '../../../../proto/language-agent/Tracing_grpc_pb';
import AuthInterceptor from '../AuthInterceptor';
import SegmentObjectAdapter from '../SegmentObjectAdapter';
import { emitter } from '../../../../lib/EventEmitter';
import Segment from '../../../../trace/context/Segment';

const logger = createLogger(__filename);

export default class TraceReportClient implements Client {
  private readonly reporterClient: TraceSegmentReportServiceClient;
  private readonly buffer: Segment[] = [];
  private timeout?: NodeJS.Timeout;
  private segmentFinishedListener: (segment: Segment) => void;

  constructor() {
    this.reporterClient = new TraceSegmentReportServiceClient(
      config.collectorAddress,
      config.secure ? grpc.credentials.createSsl() : grpc.credentials.createInsecure(),
    );

    // Store listener reference for cleanup
    this.segmentFinishedListener = (segment: Segment) => {
      // Limit buffer size to prevent memory leak during network issues
      if (this.buffer.length >= config.maxBufferSize) {
        logger.warn(
          `Trace buffer reached maximum size (${config.maxBufferSize}). ` +
          `Discarding oldest segment to prevent memory leak. ` +
          `This may indicate network connectivity issues with the collector.`
        );
        this.buffer.shift(); // Remove oldest segment
      }

      this.buffer.push(segment);
      this.timeout?.ref();
    };

    emitter.on('segment-finished', this.segmentFinishedListener);
  }

  get isConnected(): boolean {
    return this.reporterClient?.getChannel().getConnectivityState(true) === connectivityState.READY;
  }

  private reportFunction(callback?: any) {
    emitter.emit('segments-sent'); // reset limiter in SpanContext

    try {
      if (this.buffer.length === 0) {
        if (callback) callback();

        return;
      }

      const stream = this.reporterClient.collect(
        AuthInterceptor(),
        { deadline: Date.now() + config.traceTimeout },
        (error, _) => {
          if (error) {
            logger.error('Failed to report trace data', error);
          }

          if (callback) callback();
        },
      );

      try {
        for (const segment of this.buffer) {
          if (segment) {
            if (logger._isDebugEnabled) {
              logger.debug('Sending segment ', { segment });
            }

            stream.write(new SegmentObjectAdapter(segment));
          }
        }
      } finally {
        this.buffer.length = 0;
      }

      stream.end();
    } finally {
      this.timeout = setTimeout(this.reportFunction.bind(this), 1000).unref();
    }
  }

  start() {
    this.timeout = setTimeout(this.reportFunction.bind(this), 1000).unref();
  }

  flush(): Promise<any> | null {
    // This function explicitly returns null instead of a resolved Promise in case of nothing to flush so that in this
    // case passing control back to the event loop can be avoided. Even a resolved Promise will run other things in
    // the event loop when it is awaited and before it continues.

    return this.buffer.length === 0
      ? null
      : new Promise((resolve) => {
          this.reportFunction(resolve);
        });
  }

  destroy(): void {
    // Clean up event listener to prevent memory leak
    if (this.segmentFinishedListener) {
      emitter.off('segment-finished', this.segmentFinishedListener);
    }

    // Clear timeout
    if (this.timeout) {
      clearTimeout(this.timeout);
      this.timeout = undefined;
    }

    // Clear buffer
    this.buffer.length = 0;

    logger.info('TraceReportClient destroyed and resources cleaned up');
  }
}