fix: add destroy methods to prevent memory leaks in protocol clients (#129)

Co-authored-by: wolfsilver <2452450+wolfsilver@users.noreply.github.com>

Author: wolfsilver

src/agent/protocol/Protocol.ts

@@ -28,4 +28,6 @@ export default interface Protocol {
   report(): this;
 
   flush(): Promise<any> | null;
+
+  destroy?(): void;
 }

src/agent/protocol/grpc/GrpcProtocol.ts

@@ -47,4 +47,10 @@ export default class GrpcProtocol implements Protocol {
   flush(): Promise<any> | null {
     return this.traceReportClient.flush();
   }
+
+  destroy(): void {
+    // Clean up both clients to prevent memory leaks
+    this.heartbeatClient.destroy?.();
+    this.traceReportClient.destroy?.();
+  }
 }

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

@@ -23,4 +23,6 @@ export default interface Client {
   start(): void;
 
   flush(): Promise<any> | null;
+
+  destroy?(): void;
 }

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

@@ -89,4 +89,14 @@ export default class HeartbeatClient implements Client {
     logger.warn('HeartbeatClient does not need flush().');
     return null;
   }
+
+  destroy(): void {
+    // Clear heartbeat timer to prevent memory leak
+    if (this.heartbeatTimer) {
+      clearInterval(this.heartbeatTimer);
+      this.heartbeatTimer = undefined;
+    }
+
+    logger.info('HeartbeatClient destroyed and resources cleaned up');
+  }
 }

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

@@ -34,16 +34,31 @@ 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(),
     );
-    emitter.on('segment-finished', (segment) => {
+
+    // 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 {
@@ -107,4 +122,22 @@ export default class TraceReportClient implements Client {
           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');
+  }
 }

src/index.ts

@@ -72,6 +72,20 @@ class Agent {
       });
     });
   }
+
+  destroy(): void {
+    if (this.protocol === null) {
+      logger.warn('Trying to destroy() SkyWalking agent which is not started.');
+      return;
+    }
+
+    logger.info('Destroying SkyWalking agent and cleaning up resources');
+
+    // Clean up protocol resources
+    this.protocol.destroy?.();
+    this.protocol = null;
+    this.started = false;
+  }
 }
 
 export default new Agent();