Close zookeeper

service/zookeeper.cc

@@ -4,6 +4,11 @@
 
 */
 
+#include <signal.h>
+
+#include <mutex>
+#include <unordered_set>
+
 #include "soa/service/zookeeper.h"
 #include "jml/arch/timers.h"
 #include "jml/arch/backtrace.h"
@@ -13,14 +18,9 @@ using namespace std;
 
 namespace Datacratic {
 
-namespace {
+class ZookeeperConnection;
 
-struct Init {
-    Init()
-    {
-        zoo_set_debug_level(ZOO_LOG_LEVEL_ERROR);
-    }
-} init;
+namespace {
 
 void zk_callback(zhandle_t * ah, int type, int state, const char * path, void * user) {
     uintptr_t cbid = reinterpret_cast<uintptr_t>(user);
@@ -36,6 +36,37 @@ void zk_callback(zhandle_t * ah, int type, int state, const char * path, void *
     }
 }
 
+__attribute__((init_priority(101)))
+mutex closeMutex;
+
+__attribute__((init_priority(101)))
+unordered_set<ZookeeperConnection *> zookeeperConnections;
+
+struct sigaction oldAction;
+
+__attribute__((constructor))
+void
+initHandlers() {
+    zoo_set_debug_level(ZOO_LOG_LEVEL_ERROR);
+
+    sigset_t mask;
+    sigemptyset(&mask);
+    struct sigaction action = {
+         [] (int signal) {
+               unique_lock<mutex> lock{closeMutex};
+               for (auto zookeeperConnection : zookeeperConnections) {
+                  zookeeperConnection->close();
+               }
+               sigaction(SIGTERM, &oldAction, nullptr);
+               raise(SIGTERM);
+         },
+         mask,
+         0,
+         nullptr
+    };
+    sigaction(SIGTERM, &action, &oldAction);
+}
+
 } // file scope
 
 ZookeeperCallbackManager &ZookeeperCallbackManager::instance()
@@ -120,8 +151,19 @@ ZookeeperConnection()
       handle(0),
       callbackMgr_(ZookeeperCallbackManager::instance())
 {
+   unique_lock<mutex> lock{closeMutex};
+   zookeeperConnections.insert(this);
 }
-
+
+ZookeeperConnection::
+~ZookeeperConnection() {
+    close();
+    {
+        unique_lock<mutex> lock{closeMutex};
+        zookeeperConnections.erase(this);
+    }
+}
+
 std::string
 ZookeeperConnection::
 printEvent(int eventType)

service/zookeeper.h

@@ -108,7 +108,7 @@ struct ZookeeperConnection {
 
 
     ZookeeperConnection();
-    ~ZookeeperConnection() { close(); }
+    ~ZookeeperConnection();
 
     static std::string printEvent(int eventType);
 

service/zookeeper_configuration_service.cc

@@ -128,7 +128,7 @@ ZookeeperConfigurationService(std::string host,
                               std::string location,
                               int timeout)
 {
-    init(std::move(host), std::move(prefix), std::move(location));
+    init(std::move(host), std::move(prefix), std::move(location), timeout);
 }
 
 ZookeeperConfigurationService::

service/zookeeper_configuration_service.h

@@ -41,7 +41,7 @@ struct ZookeeperConfigurationService
     void init(std::string host,
               std::string prefix,
               std::string location,
-              int timeout = 5);
+              int timeout);
 
     virtual Json::Value getJson(const std::string & value,
                                 Watch watch = Watch());