add close zk fun

Author: liubaohai

cgrouparch/manager.py

@@ -7,6 +7,7 @@
 
 from pykit import threadutil
 from pykit import utfjson
+from pykit import zkutil
 from pykit.cgrouparch import account
 from pykit.cgrouparch import cgroup_manager
 from pykit.cgrouparch import communicate
@@ -85,11 +86,7 @@ def init_arch_conf(context):
         except Exception as e:
             logger.warn('failed to get arch conf from zk: %s' % repr(e))
 
-            try:
-                context['zk_client'].stop()
-            except Exception as e:
-                logger.info('failed to stop zk client: ' + repr(e))
-
+            zkutil.close_zk(context['zk_client'])
             context['zk_client'] = None
             time.sleep(10)
 
@@ -129,7 +126,7 @@ def get_cgexec_arg(cgroup_names, **argkv):
         zk_path = '%s/arch_conf' % context['zk_prefix']
         resp = zk_client.get(zk_path)
 
-        zk_client.stop()
+        zkutil.close_zk(zk_client)
 
         context['arch_conf'] = {
             'version': resp[1].version,

zktx/zktx.py

@@ -411,10 +411,7 @@ def _close(self):
         self.modifications = {}
 
         if self.owning_zk:
-            try:
-                self.zke.stop()
-            except KazooException as e:
-                logger.info(repr(e) + ' while zkclient.stop()')
+            zkutil.close_zk(self.zke)
 
     def release_all_key_locks(self):
 

zkutil/README.md

@@ -16,6 +16,7 @@
   - [zkutil.parse_kazoo_acl](#zkutilparse_kazoo_acl)
   - [zkutil.perm_to_long](#zkutilperm_to_long)
   - [zkutil.perm_to_short](#zkutilperm_to_short)
+  - [zkutil.close_zk](#zkutilclose_zk)
 - [Conditioned access methods](#conditioned-access-methods)
   - [zkutil.get_next](#zkutilget_next)
   - [zkutil.wait_absent](#zkutilwait_absent)
@@ -334,6 +335,22 @@ such as `cdrw`.
 a string of short permissions.
 
 
+##  zkutil.close_zk
+
+**syntax**:
+`zkutil.close_zk(zk)`
+
+Stop and close a zk client.
+
+**arguments**:
+
+-   `zk`:
+    a `KazooClient` or `KazooClientExt` instance, otherwise raise a `TypeError`.
+
+**return**:
+nothing
+
+
 #   Conditioned access methods
 
 A collection of helper functions those block until a condition satisfied before

zkutil/__init__.py

@@ -16,6 +16,7 @@
 from .zkutil import (
     PermTypeError,
 
+    close_zk,
     init_hierarchy,
     export_hierarchy,
     is_backward_locking,
@@ -55,6 +56,7 @@
 
     "kazoo_client_ext",
 
+    "close_zk",
     "init_hierarchy",
     "export_hierarchy",
     "is_backward_locking",

zkutil/cached_reader.py

@@ -5,8 +5,8 @@
 import threading
 
 from . import zkconf
+from . import zkutil
 from . import exceptions
-from kazoo.exceptions import KazooException
 
 logger = logging.getLogger(__name__)
 
@@ -57,10 +57,7 @@ def close(self):
 
         self.zke.remove_listener(self.conn_change_cb)
         if self.owning_zk:
-            try:
-                self.zke.stop()
-            except KazooException as e:
-                logger.exception(repr(e) + ' while close kazoo client')
+            zkutil.close_zk(self.zke)
 
     def _on_conn_change(self, state):
         logger.info('state changed: {state}'.format(state=state))

zkutil/zkacid.py

@@ -5,6 +5,7 @@
 from pykit import txutil
 
 from .zkconf import kazoo_client_ext
+from . import zkutil
 
 logger = logging.getLogger(__name__)
 
@@ -22,7 +23,4 @@ def setter(path, val, zstat):
             yield curr
     finally:
         if owning_zk:
-            try:
-                zkclient.stop()
-            except Exception as e:
-                logger.info(repr(e) + ' while stop owning zkclient')
+            zkutil.close_zk(zkclient)

zkutil/zklock.py

@@ -6,7 +6,6 @@
 import time
 
 from kazoo.client import KazooClient
-from kazoo.exceptions import KazooException
 from kazoo.exceptions import LockTimeout
 from kazoo.exceptions import NodeExistsError
 from kazoo.exceptions import NoNodeError
@@ -213,13 +212,8 @@ def close(self):
         self.zkclient.remove_listener(self.on_connection_change)
 
         if self.owning_client:
-
             logger.info('zk client is made by me, close it')
-
-            try:
-                self.zkclient.stop()
-            except KazooException as e:
-                logger.info(repr(e) + ' while stop my own client')
+            zkutil.close_zk(self.zkclient)
 
     def is_locked(self):
 

zkutil/zkutil.py

@@ -10,6 +10,7 @@
 from kazoo import security
 from kazoo.client import KazooClient
 from kazoo.exceptions import NoNodeError
+from kazoo.exceptions import KazooException
 from pykit import config
 from pykit import net
 from pykit import utfjson
@@ -47,6 +48,23 @@ class ZkPathError(Exception):
     pass
 
 
+def close_zk(zk):
+    if not isinstance(zk, KazooClient):
+        raise TypeError('expect KazooClient or KazooClientExt, but got {t}'.format(t=type(zk)))
+
+    try:
+        zk.stop()
+
+    except KazooException as e:
+        logger.exception(repr(e) + ' while stop zk client')
+
+    try:
+        zk.close()
+
+    except Exception as e:
+        logger.exception(repr(e) + ' while close zk client')
+
+
 def lock_data(node_id=None):
     # deprecated
     return lock_id(node_id=node_id)
@@ -251,7 +269,7 @@ def _init_hierarchy(hierarchy, parent_path):
             _init_hierarchy(children, path)
 
     _init_hierarchy(hierarchy, '/')
-    zkcli.stop()
+    close_zk(zkcli)
 
 
 def _make_zk_path(*paths):