From 75ada9feaa7a6713b12c8c4b8cc18735c3c1083b Mon Sep 17 00:00:00 2001
From: wusheng <wu.sheng@foxmail.com>
Date: Tue, 18 Jul 2017 21:34:30 +0800
Subject: [PATCH] Add jvm shutdown hook

---
 .../skywalking/apm/agent/core/boot/BootService.java    |  2 ++
 .../skywalking/apm/agent/core/boot/ServiceManager.java | 10 ++++++++++
 .../apm/agent/core/context/ContextManager.java         |  4 ++++
 .../java/org/skywalking/apm/agent/SkyWalkingAgent.java |  6 ++++++
 4 files changed, 22 insertions(+)

diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/boot/BootService.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/boot/BootService.java
index d3538f9152..be80b2c866 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/boot/BootService.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/boot/BootService.java
@@ -13,4 +13,6 @@ public interface BootService {
     void boot() throws Throwable;
 
     void afterBoot() throws Throwable;
+
+    void shutdown() throws Throwable;
 }
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/boot/ServiceManager.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/boot/ServiceManager.java
index 2d73389941..2c5337ef24 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/boot/ServiceManager.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/boot/ServiceManager.java
@@ -27,6 +27,16 @@ public enum ServiceManager {
         afterBoot();
     }
 
+    public void shutdown() {
+        for (BootService service : bootedServices.values()) {
+            try {
+                service.shutdown();
+            } catch (Throwable e) {
+                logger.error(e, "ServiceManager try to shutdown [{}] fail.", service.getClass().getName());
+            }
+        }
+    }
+
     private Map<Class, BootService> loadAllServices() {
         HashMap<Class, BootService> bootedServices = new HashMap<Class, BootService>();
         Iterator<BootService> serviceIterator = load().iterator();
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/ContextManager.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/ContextManager.java
index a344584645..f6347dca0e 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/ContextManager.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/ContextManager.java
@@ -162,6 +162,10 @@ public class ContextManager implements TracingContextListener, BootService, Igno
 
     }
 
+    @Override public void shutdown() throws Throwable {
+
+    }
+
     @Override
     public void afterFinished(TraceSegment traceSegment) {
         CONTEXT.remove();
diff --git a/apm-sniffer/apm-agent/src/main/java/org/skywalking/apm/agent/SkyWalkingAgent.java b/apm-sniffer/apm-agent/src/main/java/org/skywalking/apm/agent/SkyWalkingAgent.java
index ea411dd9a4..1edf9f8414 100644
--- a/apm-sniffer/apm-agent/src/main/java/org/skywalking/apm/agent/SkyWalkingAgent.java
+++ b/apm-sniffer/apm-agent/src/main/java/org/skywalking/apm/agent/SkyWalkingAgent.java
@@ -44,6 +44,12 @@ public class SkyWalkingAgent {
 
         ServiceManager.INSTANCE.boot();
 
+        Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() {
+            @Override public void run() {
+                ServiceManager.INSTANCE.shutdown();
+            }
+        }, "skywalking service shutdown thread"));
+
         new AgentBuilder.Default().type(pluginFinder.buildMatch()).transform(new AgentBuilder.Transformer() {
             @Override
             public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription,
-- 
GitLab