Add ssl support to zookeeper-statsd and fix latency handling

This adds optional SSL support to zookeeper-statsd.  This could
come in handy if we ever decide to turn off the plaintext
localhost-only port.

This also corrects the type handling for the latency value, which
can be a floating point.

Change-Id: Id39fc8bd924eda528723c40d2e7e24993a60d6a5
This commit is contained in:
James E. Blair 2022-11-09 13:44:28 -08:00
parent 2c72c8b9e3
commit a7026aba8a

View File

@ -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()