From 793fffdc20da9ad50e0a190c7936995896bfd1f9 Mon Sep 17 00:00:00 2001
From: Sahil Sabharwal <ssabharwal@juniper.net>
Date: Thu, 1 Dec 2016 14:23:13 -0800
Subject: [PATCH] Added Zookeeper timeout config knob in ST

Config knob "zk_timeout" for Zookeeper timeout is added in Schema Transformer.
ZookeeperClient now takes in arguemet for zk_timeout for which default has been
set to 400.

Change-Id: Ia1a5569292f4957dba4b3f64aaee997d7db9f9da
Closes-Bug: 1643846
---
 src/config/common/zkclient.py           | 4 ++--
 src/config/schema-transformer/to_bgp.py | 7 ++++++-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/config/common/zkclient.py b/src/config/common/zkclient.py
index 56fe706cd15..4965dfba6e4 100644
--- a/src/config/common/zkclient.py
+++ b/src/config/common/zkclient.py
@@ -220,7 +220,7 @@ def delete_all(cls, zookeeper_client, path):
 
 class ZookeeperClient(object):
 
-    def __init__(self, module, server_list, logging_fn=None):
+    def __init__(self, module, server_list, logging_fn=None, zk_timeout=400):
         # logging
         logger = logging.getLogger(module)
         logger.setLevel(logging.DEBUG)
@@ -245,7 +245,7 @@ def __init__(self, module, server_list, logging_fn=None):
                                  sleep_func=gevent.sleep)
         self._zk_client = kazoo.client.KazooClient(
                 server_list,
-                timeout=400,
+                timeout=zk_timeout,
                 handler=kazoo.handlers.gevent.SequentialGeventHandler(),
                 logger=logger,
                 connection_retry=self._retry,
diff --git a/src/config/schema-transformer/to_bgp.py b/src/config/schema-transformer/to_bgp.py
index a8f8fb502a7..57b0780ff81 100644
--- a/src/config/schema-transformer/to_bgp.py
+++ b/src/config/schema-transformer/to_bgp.py
@@ -641,6 +641,7 @@ def parse_args(args_str):
                          --use_syslog
                          --syslog_facility LOG_USER
                          --cluster_id <testbed-name>
+                         --zk_timeout 400
                          [--reset_config]
     '''
 
@@ -687,6 +688,7 @@ def parse_args(args_str):
         'kombu_ssl_keyfile': '',
         'kombu_ssl_certfile': '',
         'kombu_ssl_ca_certs': '',
+        'zk_timeout': 400,
     }
     secopts = {
         'use_certs': False,
@@ -803,6 +805,8 @@ def parse_args(args_str):
                         help="Start port for bgp-as-a-service proxy")
     parser.add_argument("--bgpaas_port_end", type=int,
                         help="End port for bgp-as-a-service proxy")
+    parser.add_argument("--zk_timeout",
+                        help="Timeout for ZookeeperClient")
 
     args = parser.parse_args(remaining_argv)
     if type(args.cassandra_server_list) is str:
@@ -862,7 +866,8 @@ def main(args_str=None):
     else:
         client_pfx = ''
         zk_path_pfx = ''
-    _zookeeper_client = ZookeeperClient(client_pfx+"schema", args.zk_server_ip)
+    _zookeeper_client = ZookeeperClient(client_pfx+"schema", args.zk_server_ip,
+                                        zk_timeout =args.zk_timeout)
     _zookeeper_client.master_election(zk_path_pfx + "/schema-transformer",
                                       os.getpid(), run_schema_transformer,
                                       args)