#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
  varnishstatsd
  ~~~~~~~~~~~~~
  Report backend response times and request counts aggregated by status.

  Usage: varnishstatsd [--varnish-name NAME] [--statsd-server SERVER]
                       [--key-prefix PREFIX]

    --varnish-name NAME     varnish name
    --statsd-server SERVER  statsd server
    --key-prefix PREFIX     metric key prefix

  If no statsd server is specified, prints stats to stdout instead.

  Copyright 2015 Ori Livneh <ori@wikimedia.org>

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.

"""
import argparse
import collections
import io
import socket
import urlparse
from varnishlog import varnishlog


METRIC_FORMAT = (
    '%(key_prefix)s%(backend)s.%(method)s:%(ttfb)d|ms\n'
    '%(key_prefix)s%(backend)s.%(status)s:1|c\n'
)

UDP_MTU_BYTES = 1472

vsl_args = [
    ('i', 'Backend'),
    ('i', 'BackendXID'),
    ('i', 'ReqEnd'),
    ('i', 'RxStatus'),
    ('i', 'TxRequest'),
]

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)


def parse_statsd_server_string(server_string):
    parsed = urlparse.urlparse('//' + server_string)
    return parsed.hostname, parsed.port or 8125


def parse_key_prefix(key_prefix):
    return key_prefix.strip('.') + '.'


ap = argparse.ArgumentParser(
    description='Varnish backend response time metric logger',
    epilog='If no statsd server is specified, prints stats to stdout instead.'
)
ap.add_argument('--varnish-name', help='varnish name')
ap.add_argument('--statsd-server', help='statsd server',
                type=parse_statsd_server_string, default=None)
ap.add_argument('--key-prefix', help='metric key prefix',
                default='varnish.backends',
                type=parse_key_prefix)
args = ap.parse_args()

if args.varnish_name:
    vsl_args.append(('n', args.varnish_name))

# Maximum number of log records to keep in memory.
MAX_BACKLOG = 100

# Map of transaction ID: request method ('GET', 'POST', etc.)
tids = collections.OrderedDict()

# Map of backend XID: transaction ID
xids = collections.OrderedDict()

# Map of transaction ID: backend name
backends = collections.OrderedDict()

# Map of transaction ID: response status code (200, 404, etc.)
statuses = collections.OrderedDict()

buf = io.BytesIO()


def log_value_err(callback):
    def wrapper(transaction_id, tag, record, remote_party):
        try:
            return callback(transaction_id, tag, record, remote_party)
        except ValueError as err:
            print locals()
            return

    return wrapper


@log_value_err
def vsl_callback(transaction_id, tag, record, remote_party):
    global buf

    if tag == 'Backend':
        # Associate the transaction ID with the backend name.
        tid, _, backend = record.split()
        backends[int(tid)] = backend
        if len(backends) > MAX_BACKLOG:
            backends.popitem(last=False)

    elif tag == 'BackendXID':
        # Associate the backend XID with its transaction ID.
        xid = int(record) & 0xffffffff
        xids[xid] = transaction_id
        if len(xids) > MAX_BACKLOG:
            xids.popitem(last=False)

    elif tag == 'TxRequest' and remote_party == 'backend':
        # Associate the transaction ID with the request method.
        tids[transaction_id] = record
        if len(tids) > MAX_BACKLOG:
            tids.popitem(last=False)

    elif tag == 'RxStatus' and remote_party == 'backend':
        # Associate the transaction ID with the response status code.
        if record[0] in '12345':
            statuses[transaction_id] = record[0] + 'xx'
        if len(statuses) > MAX_BACKLOG:
            statuses.popitem(last=False)

    elif tag == 'ReqEnd':
        # Collate data and emit metric.
        parts = record.split()
        xid = int(parts[0])

        try:
            tid = xids.pop(xid)
            fields = {
                'key_prefix': args.key_prefix,
                'method': tids.pop(tid),
                'backend': backends.pop(tid),
                'status': statuses.pop(tid),
                'ttfb': round(1000 * float(parts[4])),
            }
        except KeyError:
            return 0

        metric_string = (METRIC_FORMAT % fields).encode('utf-8')
        if buf.tell() + len(metric_string) >= UDP_MTU_BYTES:
            buf.seek(io.SEEK_SET)
            if args.statsd_server:
                sock.sendto(buf.read(), args.statsd_server)
            else:
                print(buf.read().decode('utf-8', errors='replace').rstrip())
            buf = io.BytesIO()
        buf.write(metric_string)

    return 0

varnishlog(vsl_args, vsl_callback)
