From 4f022b30558ea18a2f9fed5518828f069ec65601 Mon Sep 17 00:00:00 2001
From: Jonas Dohse <jonas@mbr-targeting.com>
Date: Tue, 5 Nov 2013 10:01:57 +0000
Subject: [PATCH 1/2] Forward timeout to initialization method

---
 service/zookeeper_configuration_service.cc | 2 +-
 service/zookeeper_configuration_service.h  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/service/zookeeper_configuration_service.cc b/service/zookeeper_configuration_service.cc
index a2e07bbf..963fb4f7 100644
--- a/service/zookeeper_configuration_service.cc
+++ b/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::
diff --git a/service/zookeeper_configuration_service.h b/service/zookeeper_configuration_service.h
index 30c8e71a..4fb5421e 100644
--- a/service/zookeeper_configuration_service.h
+++ b/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());

From df64f82a290bbce3a6311047ee0086dec4b716da Mon Sep 17 00:00:00 2001
From: Jonas Dohse <jonas@mbr-targeting.com>
Date: Mon, 4 Nov 2013 13:17:20 +0000
Subject: [PATCH 2/2] Close zookeeper handles on SIGTERM

---
 service/zookeeper.cc | 58 ++++++++++++++++++++++++++++++++++++++------
 service/zookeeper.h  |  2 +-
 2 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/service/zookeeper.cc b/service/zookeeper.cc
index b8c9fa03..c7851c29 100644
--- a/service/zookeeper.cc
+++ b/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)
diff --git a/service/zookeeper.h b/service/zookeeper.h
index ea1f06bd..9bfbf2e8 100644
--- a/service/zookeeper.h
+++ b/service/zookeeper.h
@@ -108,7 +108,7 @@ struct ZookeeperConnection {
 
 
     ZookeeperConnection();
-    ~ZookeeperConnection() { close(); }
+    ~ZookeeperConnection();
 
     static std::string printEvent(int eventType);