BATCAVE/ffstatus/graphite.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

from __future__ import print_function
import logging
import socket
import time
import StringIO

class GraphitePush:
    dont_send = False

    prefix = 'ffpb.'
    target_host = None
    target_port = 2003

    whitelist = None

    def __init__(self, host, port=2003):
        self.target_host = host
        self.target_port = port

    def __str__(self):
        return 'Graphite at [{0}]:{1} (prefix=\'{2}\', whitelist={3})'.format(
            self.target_host, self.target_port,
            self.prefix, self.whitelist)

    def __print_graphite_line(self, output, nodeid, item, value, timestamp):
        print(
            self.prefix, 'nodes.', nodeid, '.', item, ' ', value, ' ', timestamp,
            sep='', file=output)

    def __do_send(self, data):
        sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
        sock.connect((self.target_host, self.target_port))
        sock.sendall(data)
        sock.shutdown(socket.SHUT_WR)
        sock.close()

    def push(self, data, timestamp=None):
        if timestamp is None:
            timestamp = time.time()
        timestamp = int(timestamp)

        output = StringIO.StringIO()
        whitelist = None
        if self.whitelist is not None and len(self.whitelist) > 0:
            whitelist = [x for x in self.whitelist]

        for nodeid in data:
            if whitelist is not None and nodeid not in whitelist:
                #logging.debug(
                #    'Graphite output skips node \'%s\' (not in whitelist).',
                #    nodeid)
                continue

            nodeinfo = data[nodeid]

            for item in ['uptime']:
                if item in nodeinfo:
                    self.__print_graphite_line(
                        output, nodeid, item, nodeinfo[item], timestamp)

            traffic = nodeinfo.get('statistics', {}).get('traffic')
            if traffic is not None:
                for item in ['rxbytes', 'txbytes']:
                    self.__print_graphite_line(
                        output, nodeid, item, traffic[item], timestamp)

        all_output = output.getvalue()

        if not self.dont_send:
            self.__do_send(all_output)

        output.close()

        return all_output

    def handle_metric(self, source, key, value, ts):
        if self.whitelist is not None and len(self.whitelist) > 0:
            valid_prefixes = ['nodes.' + x for x in self.whitelist]
            if len([x for x in valid_prefixes if key.startswith(x)]) == 0:
                # not in whitelist
                return

        data = '{key} {value} {ts}\n'.format(
            key=self.prefix + str(key).replace(' ', '_'),
            value=float(value),
            ts=int(ts))
        self.__do_send(data)