From 8e89970c70ec145254986c15376b1f450e2f4e76 Mon Sep 17 00:00:00 2001
From: wolfsilver <2452450+wolfsilver@users.noreply.github.com>
Date: Wed, 28 Jan 2026 11:25:53 +0800
Subject: [PATCH] fix: add destroy methods to prevent memory leaks in protocol
 clients

---
 src/agent/protocol/Protocol.ts                |  2 +
 src/agent/protocol/grpc/GrpcProtocol.ts       |  6 +++
 src/agent/protocol/grpc/clients/Client.ts     |  2 +
 .../protocol/grpc/clients/HeartbeatClient.ts  | 10 +++++
 .../grpc/clients/TraceReportClient.ts         | 37 ++++++++++++++++++-
 src/index.ts                                  | 14 +++++++
 6 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/src/agent/protocol/Protocol.ts b/src/agent/protocol/Protocol.ts
index 842f3bc..062318a 100644
--- a/src/agent/protocol/Protocol.ts
+++ b/src/agent/protocol/Protocol.ts
@@ -28,4 +28,6 @@ export default interface Protocol {
   report(): this;
 
   flush(): Promise<any> | null;
+
+  destroy?(): void;
 }
diff --git a/src/agent/protocol/grpc/GrpcProtocol.ts b/src/agent/protocol/grpc/GrpcProtocol.ts
index 08ffa63..f151141 100644
--- a/src/agent/protocol/grpc/GrpcProtocol.ts
+++ b/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?.();
+  }
 }
diff --git a/src/agent/protocol/grpc/clients/Client.ts b/src/agent/protocol/grpc/clients/Client.ts
index e9916ab..8d8fd06 100644
--- a/src/agent/protocol/grpc/clients/Client.ts
+++ b/src/agent/protocol/grpc/clients/Client.ts
@@ -23,4 +23,6 @@ export default interface Client {
   start(): void;
 
   flush(): Promise<any> | null;
+
+  destroy?(): void;
 }
diff --git a/src/agent/protocol/grpc/clients/HeartbeatClient.ts b/src/agent/protocol/grpc/clients/HeartbeatClient.ts
index b92d6e3..91cda52 100755
--- a/src/agent/protocol/grpc/clients/HeartbeatClient.ts
+++ b/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');
+  }
 }
diff --git a/src/agent/protocol/grpc/clients/TraceReportClient.ts b/src/agent/protocol/grpc/clients/TraceReportClient.ts
index 3f6f102..8352604 100755
--- a/src/agent/protocol/grpc/clients/TraceReportClient.ts
+++ b/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');
+  }
 }
diff --git a/src/index.ts b/src/index.ts
index a228620..ae2c494 100644
--- a/src/index.ts
+++ b/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();