Fix DubboShutdownHook Memory Leak (#2922)

apache · Dec 10, 2018 · 833ba03 · 833ba03
1 parent c2a4c8d
commit 833ba03
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 4 deletions.
diff --git a/dubbo-config/dubbo-config-api/src/main/java/org/apache/dubbo/config/AbstractConfig.java b/dubbo-config/dubbo-config-api/src/main/java/org/apache/dubbo/config/AbstractConfig.java
@@ -74,7 +74,7 @@ public abstract class AbstractConfig implements Serializable {
         legacyProperties.put("dubbo.service.url", "dubbo.service.address");
 
         // this is only for compatibility
-        Runtime.getRuntime().addShutdownHook(DubboShutdownHook.getDubboShutdownHook());
+        DubboShutdownHook.getDubboShutdownHook().register();
     }
 
     protected String id;

diff --git a/dubbo-config/dubbo-config-api/src/main/java/org/apache/dubbo/config/DubboShutdownHook.java b/dubbo-config/dubbo-config-api/src/main/java/org/apache/dubbo/config/DubboShutdownHook.java
@@ -40,14 +40,17 @@ public static DubboShutdownHook getDubboShutdownHook() {
         return dubboShutdownHook;
     }
 
+    /**
+     * Has it already been registered or not?
+     */
+    private final AtomicBoolean registered = new AtomicBoolean(false);
     /**
      * Has it already been destroyed or not?
      */
-    private final AtomicBoolean destroyed;
+    private final AtomicBoolean destroyed= new AtomicBoolean(false);
 
     private DubboShutdownHook(String name) {
         super(name);
-        this.destroyed = new AtomicBoolean(false);
     }
 
     @Override
@@ -58,13 +61,33 @@ public void run() {
         doDestroy();
     }
 
+    /**
+     * Register the ShutdownHook
+     */
+    public void register() {
+        if (!registered.get() && registered.compareAndSet(false, true)) {
+            Runtime.getRuntime().addShutdownHook(getDubboShutdownHook());
+        }
+    }
+
+    /**
+     * Unregister the ShutdownHook
+     */
+    public void unregister() {
+        if (registered.get() && registered.compareAndSet(true, false)) {
+            Runtime.getRuntime().removeShutdownHook(getDubboShutdownHook());
+        }
+    }
+
     /**
      * Destroy all the resources, including registries and protocols.
      */
     public void doDestroy() {
         if (!destroyed.compareAndSet(false, true)) {
             return;
         }
+        // unregister the shutdownHook
+        unregister();
         // destroy all the registries
         AbstractRegistryFactory.destroyAll();
         // destroy all the protocols
@@ -88,4 +111,5 @@ private void destroyProtocols() {
         }
     }
 
+
 }
diff --git a/...spring/src/main/java/org/apache/dubbo/config/spring/extension/SpringExtensionFactory.java b/...spring/src/main/java/org/apache/dubbo/config/spring/extension/SpringExtensionFactory.java
@@ -44,7 +44,7 @@ public class SpringExtensionFactory implements ExtensionFactory {
 
     public static void addApplicationContext(ApplicationContext context) {
         contexts.add(context);
-        Runtime.getRuntime().removeShutdownHook(DubboShutdownHook.getDubboShutdownHook());
+        DubboShutdownHook.getDubboShutdownHook().unregister();
         BeanFactoryUtils.addApplicationListener(context, shutdownHookListener);
     }