Improve graceful shutdown

dubbo-bootstrap/src/main/java/org/apache/dubbo/bootstrap/DubboBootstrap.java

@@ -16,66 +16,65 @@
  */
 package org.apache.dubbo.bootstrap;
 
-import com.alibaba.dubbo.common.extension.ExtensionLoader;
-import com.alibaba.dubbo.common.logger.Logger;
-import com.alibaba.dubbo.common.logger.LoggerFactory;
+import com.alibaba.dubbo.config.DubboShutdownHook;
 import com.alibaba.dubbo.config.ServiceConfig;
-import com.alibaba.dubbo.registry.support.AbstractRegistryFactory;
-import com.alibaba.dubbo.rpc.Protocol;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * A bootstrap class to easily start and stop Dubbo via programmatic API.
  * The bootstrap class will be responsible to cleanup the resources during stop.
  */
 public class DubboBootstrap {
 
-    private static final Logger logger = LoggerFactory.getLogger(DubboBootstrap.class);
-
     /**
      * The list of ServiceConfig
      */
     private List<ServiceConfig> serviceConfigList;
 
     /**
-     * Has it already been destroyed or not?
+     * Whether register the shutdown hook during start?
      */
-    private final AtomicBoolean destroyed;
+    private final boolean registerShutdownHookOnStart;
 
     /**
      * The shutdown hook used when Dubbo is running under embedded environment
      */
-    private Thread shutdownHook;
+    private DubboShutdownHook shutdownHook;
 
     public DubboBootstrap() {
+        this(true, DubboShutdownHook.getDubboShutdownHook());
+    }
+
+    public DubboBootstrap(boolean registerShutdownHookOnStart) {
+        this(registerShutdownHookOnStart, DubboShutdownHook.getDubboShutdownHook());
+    }
+
+    public DubboBootstrap(boolean registerShutdownHookOnStart, DubboShutdownHook shutdownHook) {
         this.serviceConfigList = new ArrayList<ServiceConfig>();
-        this.destroyed = new AtomicBoolean(false);
-        this.shutdownHook = new Thread(new Runnable() {
-            @Override
-            public void run() {
-                if (logger.isInfoEnabled()) {
-                    logger.info("Run shutdown hook now.");
-                }
-                destroy();
-            }
-        }, "DubboShutdownHook");
+        this.shutdownHook = shutdownHook;
+        this.registerShutdownHookOnStart = registerShutdownHookOnStart;
     }
 
     /**
      * Register service config to bootstrap, which will be called during {@link DubboBootstrap#stop()}
      * @param serviceConfig the service
      * @return the bootstrap instance
      */
-    public DubboBootstrap regsiterServiceConfig(ServiceConfig serviceConfig) {
+    public DubboBootstrap registerServiceConfig(ServiceConfig serviceConfig) {
         serviceConfigList.add(serviceConfig);
         return this;
     }
 
     public void start() {
-        registerShutdownHook();
+        if (registerShutdownHookOnStart) {
+            registerShutdownHook();
+        } else {
+            // DubboShutdown hook has been registered in AbstractConfig,
+            // we need to remove it explicitly
+            removeShutdownHook();
+        }
         for (ServiceConfig serviceConfig: serviceConfigList) {
             serviceConfig.export();
         }
@@ -85,8 +84,10 @@ public void stop() {
         for (ServiceConfig serviceConfig: serviceConfigList) {
             serviceConfig.unexport();
         }
-        destroy();
-        removeShutdownHook();
+        shutdownHook.destroyAll();
+        if (registerShutdownHookOnStart) {
+            removeShutdownHook();
+        }
     }
 
     /**
@@ -107,27 +108,4 @@ public void removeShutdownHook() {
             // ignore - VM is already shutting down
         }
     }
-
-    /**
-     * Destroy all the resources, including registries and protocols.
-     */
-    private void destroy() {
-        if (!destroyed.compareAndSet(false, true)) {
-            return;
-        }
-        // destroy all the registries
-        AbstractRegistryFactory.destroyAll();
-        // destroy all the protocols
-        ExtensionLoader<Protocol> loader = ExtensionLoader.getExtensionLoader(Protocol.class);
-        for (String protocolName : loader.getLoadedExtensions()) {
-            try {
-                Protocol protocol = loader.getLoadedExtension(protocolName);
-                if (protocol != null) {
-                    protocol.destroy();
-                }
-            } catch (Throwable t) {
-                logger.warn(t.getMessage(), t);
-            }
-        }
-    }
 }

dubbo-config/dubbo-config-api/src/main/java/com/alibaba/dubbo/config/AbstractConfig.java

@@ -71,6 +71,9 @@ public abstract class AbstractConfig implements Serializable {
         legacyProperties.put("dubbo.consumer.retries", "dubbo.service.max.retry.providers");
         legacyProperties.put("dubbo.consumer.check", "dubbo.service.allow.no.provider");
         legacyProperties.put("dubbo.service.url", "dubbo.service.address");
+
+        // this is only for compatibility
+        Runtime.getRuntime().addShutdownHook(DubboShutdownHook.getDubboShutdownHook());
     }
 
     protected String id;

dubbo-config/dubbo-config-api/src/main/java/com/alibaba/dubbo/config/DubboShutdownHook.java

@@ -0,0 +1,85 @@
+/*
+ * 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.
+ */
+package com.alibaba.dubbo.config;
+
+import com.alibaba.dubbo.common.extension.ExtensionLoader;
+import com.alibaba.dubbo.common.logger.Logger;
+import com.alibaba.dubbo.common.logger.LoggerFactory;
+import com.alibaba.dubbo.registry.support.AbstractRegistryFactory;
+import com.alibaba.dubbo.rpc.Protocol;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * The shutdown hook thread to do the clean up stuff.
+ * This is a singleton in order to ensure there is only one shutdown hook registered.
+ * Because {@link ApplicationShutdownHooks} use {@link java.util.IdentityHashMap}
+ * to store the shutdown hooks.
+ */
+public class DubboShutdownHook extends Thread {
+
+    private static final Logger logger = LoggerFactory.getLogger(DubboShutdownHook.class);
+
+    private static final DubboShutdownHook dubboShutdownHook = new DubboShutdownHook("DubboShutdownHook");
+
+    public static DubboShutdownHook getDubboShutdownHook() {
+        return dubboShutdownHook;
+    }
+
+    /**
+     * Has it already been destroyed or not?
+     */
+    private final AtomicBoolean destroyed;
+
+    private DubboShutdownHook(String name) {
+        super(name);
+        this.destroyed = new AtomicBoolean(false);
+    }
+
+    @Override
+    public void run() {
+        if (logger.isInfoEnabled()) {
+            logger.info("Run shutdown hook now.");
+        }
+        destroyAll();
+    }
+
+    /**
+     * Destroy all the resources, including registries and protocols.
+     */
+    public void destroyAll() {
+        if (!destroyed.compareAndSet(false, true)) {
+            return;
+        }
+        // destroy all the registries
+        AbstractRegistryFactory.destroyAll();
+        // destroy all the protocols
+        ExtensionLoader<Protocol> loader = ExtensionLoader.getExtensionLoader(Protocol.class);
+        for (String protocolName : loader.getLoadedExtensions()) {
+            try {
+                Protocol protocol = loader.getLoadedExtension(protocolName);
+                if (protocol != null) {
+                    protocol.destroy();
+                }
+            } catch (Throwable t) {
+                logger.warn(t.getMessage(), t);
+            }
+        }
+    }
+
+
+}

dubbo-config/dubbo-config-api/src/main/java/com/alibaba/dubbo/config/ProtocolConfig.java

@@ -460,4 +460,12 @@ public void destory() {
         }
     }
 
+    /**
+     * Just for compatibility.
+     * It should be deleted in the next major version, say 2.7.x.
+     */
+    @Deprecated
+    public static void destroyAll() {
+        DubboShutdownHook.getDubboShutdownHook().destroyAll();
+    }
 }

dubbo-config/dubbo-config-spring/src/main/java/com/alibaba/dubbo/config/spring/initializer/DubboApplicationListener.java

@@ -31,7 +31,11 @@ public class DubboApplicationListener implements ApplicationListener<Application
     private DubboBootstrap dubboBootstrap;
 
     public DubboApplicationListener() {
-        dubboBootstrap = new DubboBootstrap();
+        dubboBootstrap = new DubboBootstrap(false);
+    }
+
+    public DubboApplicationListener(DubboBootstrap dubboBootstrap) {
+        this.dubboBootstrap = dubboBootstrap;
     }
 
     @Override

dubbo-config/dubbo-config-spring/src/test/java/com/alibaba/dubbo/config/spring/initializer/DubboApplicationListenerTest.java

@@ -0,0 +1,59 @@
+/*
+ * 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.
+ */
+package com.alibaba.dubbo.config.spring.initializer;
+
+import com.alibaba.dubbo.config.DubboShutdownHook;
+import org.apache.dubbo.bootstrap.DubboBootstrap;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.springframework.context.support.ClassPathXmlApplicationContext;
+
+public class DubboApplicationListenerTest {
+
+    @Test
+    public void testTwoShutdownHook() {
+        DubboShutdownHook spyHook = Mockito.spy(DubboShutdownHook.getDubboShutdownHook());
+        ClassPathXmlApplicationContext applicationContext = getApplicationContext(spyHook, true);
+        applicationContext.refresh();
+        applicationContext.close();
+        // shutdown hook can't be verified, because it will executed after main thread has finished.
+        // so we can only verify it by manually run it.
+        try {
+            spyHook.start();
+            spyHook.join();
+        } catch (InterruptedException e) {
+            e.printStackTrace();
+        }
+        Mockito.verify(spyHook, Mockito.times(2)).destroyAll();
+    }
+
+    @Test
+    public void testOneShutdownHook() {
+        DubboShutdownHook spyHook = Mockito.spy(DubboShutdownHook.getDubboShutdownHook());
+        ClassPathXmlApplicationContext applicationContext = getApplicationContext(spyHook, false);
+        applicationContext.refresh();
+        applicationContext.close();
+        Mockito.verify(spyHook, Mockito.times(1)).destroyAll();
+    }
+
+    private ClassPathXmlApplicationContext getApplicationContext(DubboShutdownHook hook, boolean registerHook) {
+        DubboBootstrap bootstrap = new DubboBootstrap(registerHook, hook);
+        ClassPathXmlApplicationContext applicationContext = new ClassPathXmlApplicationContext();
+        applicationContext.addApplicationListener(new DubboApplicationListener(bootstrap));
+        return applicationContext;
+    }
+}

dubbo-container/dubbo-container-spring/pom.xml

@@ -39,5 +39,10 @@
             <groupId>org.springframework</groupId>
             <artifactId>spring-context</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.alibaba</groupId>
+            <artifactId>dubbo-config-spring</artifactId>
+            <version>${project.parent.version}</version>
+        </dependency>
     </dependencies>
 </project>

dubbo-container/dubbo-container-spring/src/main/java/com/alibaba/dubbo/container/spring/SpringContainer.java

@@ -19,6 +19,7 @@
 import com.alibaba.dubbo.common.logger.Logger;
 import com.alibaba.dubbo.common.logger.LoggerFactory;
 import com.alibaba.dubbo.common.utils.ConfigUtils;
+import com.alibaba.dubbo.config.spring.initializer.DubboApplicationListener;
 import com.alibaba.dubbo.container.Container;
 
 import org.springframework.context.support.ClassPathXmlApplicationContext;
@@ -44,6 +45,7 @@ public void start() {
             configPath = DEFAULT_SPRING_CONFIG;
         }
         context = new ClassPathXmlApplicationContext(configPath.split("[,\\s]+"));
+        context.addApplicationListener(new DubboApplicationListener());
         context.start();
     }