Kill old haproxy process Raw

marathonlb-kill-old-process.py

#!/usr/bin/env python2

# pip install pid, psutil, raven

import argparse
import datetime
import logging
import logging.handlers
import sys
import time
from collections import Counter
from itertools import ifilter as filter
from itertools import ifilterfalse as filterfalse

import psutil
from pid.decorator import pidfile

APP_NAME = 'marathonlb-kill-old-process'
SENTRY_DSN = ''


def configure_logger():
    logger = logging.getLogger(APP_NAME)
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    ch = logging.StreamHandler()
    fh = logging.handlers.RotatingFileHandler(
        '/var/log/{}/{}.log'.format(APP_NAME, APP_NAME),
        maxBytes=10 * 1024 * 1024,
        backupCount=2)
    ch.setFormatter(formatter)
    fh.setFormatter(formatter)
    logger.addHandler(ch)
    logger.addHandler(fh)
    return logger


logger = configure_logger()


def is_haproxy(p):
    try:
        return p.name() == 'haproxy'
    except Exception:
        return False


def can_kill(p):
    try:
        status = set(c.status for c in p.connections())
        return psutil.CONN_LISTEN not in status
    except Exception:
        return False


def is_old_enough(p):
    try:
        create_time = datetime.datetime.fromtimestamp(p.create_time())
        now = datetime.datetime.now()
        delta = now - create_time
        five_minutes = 60 * 5
        return delta.seconds > five_minutes
    except Exception:
        return False


def process_to_kill():
    procs = filter(is_old_enough,
                   filter(can_kill, filter(is_haproxy, psutil.process_iter())))
    return procs


def wait(sleep_time=2):
    logger.info('Waiting %d seconds', sleep_time)
    time.sleep(sleep_time)


def get_listen_ports():
    try:
        procs = list(
            filterfalse(can_kill, filter(is_haproxy, psutil.process_iter())))
        ports = set(c.laddr[1] for p in procs for c in p.connections()
                    if c.status == psutil.CONN_LISTEN for p in procs)
        return ports
    except Exception:
        return []


def count_conn_of_listen_ports(p, listen_ports):
    if not listen_ports:
        return {'???': '???'}
    c = Counter(c.laddr[1] for c in p.connections()
                if c.laddr[1] in listen_ports)
    return dict(c.viewitems())


def kill(p, dry_run, listen_ports):
    try:
        name = p.name()
        pid = p.pid
        count_conn = count_conn_of_listen_ports(p, listen_ports)
        if not dry_run:
            p.kill()
        logger.info('killing %s of pid %d with lports %s', name, pid,
                    count_conn)
        return True
    except Exception:
        return False


@pidfile(APP_NAME)
def main(dry_run=False, sleep_time=2):
    logger.info('Seek and Destroy \m/')
    count = 0
    listen_ports = get_listen_ports()
    for proc in process_to_kill():
        wait(sleep_time)
        if kill(proc, dry_run, listen_ports):
            count += 1
    logger.info('Finished killing %d process', count)


def parse_args():
    parser = argparse.ArgumentParser(
        description='Remove all old haproxy process')
    parser.add_argument(
        '--dry-run',
        help='Just print the actions and set sleep time to 0',
        required=False,
        action='store_true')
    parser.add_argument(
        '--sleep-time',
        help='Sleep time between each kill',
        required=False,
        type=int,
        default=2)
    conf = vars(parser.parse_args(sys.argv[1:]))
    if conf.get('dry_run', False):
        conf['sleep_time'] = 0
    return conf


if __name__ == '__main__':
    from raven import Client
    client = Client(SENTRY_DSN)

    try:
        conf = parse_args()
        main(**conf)
    except Exception:
        client.captureException()