@ -21,6 +21,8 @@ import logging
import re
import socket
import time
import os
import ssl
from statsd . defaults . env import statsd
@ -48,14 +50,23 @@ COUNTERS = [
class Socket :
def __init__ ( self , host , port ):
def __init__ ( self , host , port , ca_cert , client_cert , client_key ):
self . host = host
self . port = port
self . ca_cert = ca_cert
self . client_cert = client_cert
self . client_key = client_key
self . socket = None
def open ( self ) :
s = socket . socket ( socket . AF_INET , socket . SOCK_STREAM )
s . settimeout ( 5 )
if self . client_key :
context = ssl . SSLContext ( ssl . PROTOCOL_TLS_CLIENT )
context . load_verify_locations ( self . ca_cert )
context . load_cert_chain ( self . client_cert , self . client_key )
context . check_hostname = False
s = context . wrap_socket ( s , server_hostname = self . host )
s . connect ( ( self . host , self . port ) )
self . socket = s
@ -69,8 +80,13 @@ class Socket:
class ZooKeeperStats :
def __init__ ( self , host , port = 2181 ) :
self . socket = Socket ( host , port )
def __init__ ( self , host , port = None ,
ca_cert = None , client_cert = None , client_key = None ) :
if client_key :
port = port or 2281
else :
port = port or 2181
self . socket = Socket ( host , port , ca_cert , client_cert , client_key )
# The hostname to use when reporting stats (e.g., zk01)
if host in ( ' localhost ' , ' 127.0.0.1 ' , ' ::1 ' ) :
self . hostname = socket . gethostname ( )
@ -108,7 +124,11 @@ class ZooKeeperStats:
base = ' zk. %s . ' % ( self . hostname , )
for key in GAUGES :
try :
value = int ( stats . get ( key , 0 ) )
value = stats . get ( key , ' 0 ' )
if ' . ' in value :
value = float ( value )
else :
value = int ( value )
pipe . gauge ( base + key , value )
except Exception :
self . log . exception ( " Unable to process %s " , key )
@ -137,6 +157,11 @@ class ZooKeeperStats:
self . reportStats ( stats )
ca_cert = os . environ . get ( " ZK_CA_CERT " )
client_cert = os . environ . get ( " ZK_CLIENT_CERT " )
client_key = os . environ . get ( " ZK_CLIENT_KEY " )
logging . basicConfig ( level = logging . DEBUG )
p = ZooKeeperStats ( ' localhost ' )
p = ZooKeeperStats ( ' localhost ' , ca_cert = ca_cert ,
client_cert = client_cert , client_key = client_key )
p . run ( )