tasktiger remove orphan tasks

remove_orphan_tasks.py

# -*- coding: utf-8 -*-
from __future__ import unicode_literals, print_function, division
import sys
from redis import Redis


def check():
    r = Redis()
    task_keys = set()
    for key in r.scan_iter(match='t:task:*'):
        if key.endswith('executions'):
            key = key.replace(':executions', '')
        task_keys.add(key)
    print('total task keys: {}'.format(len(task_keys)))

    in_queue_task_keys = set()
    for state in {'queued', 'active', 'scheduled', 'error'}:
        for queue_name in r.smembers('t:{}'.format(state)):
            for key in r.zrange('t:{}:{}'.format(state, queue_name), 0, -1):
                in_queue_task_keys.add('t:task:{}'.format(key))
    print('total in queue task keys: {}'.format(len(in_queue_task_keys)))

    orphan_task_keys = task_keys - in_queue_task_keys
    if not orphan_task_keys:
        print('no orphan tasks')
    else:
        with open('orphan_tasks', mode='wb+') as f:
            tasks = r.mget(*orphan_task_keys)
            for task in tasks:
                f.write(b'{}\n'.format(task))
        print('{} orphan task saved to file: orphan_tasks'.format(len(orphan_task_keys)))
        with open('orphan_task_keys', mode='wb+') as f:
            for key in orphan_task_keys:
                f.write(b'{}\n'.format(key))
        print('orphan task keys saved to file: orphan_task_keys')


def remove(key_file):
    with open(key_file) as f:
        lines = f.readlines()

    lines = [line.strip() for line in lines if line.strip()]
    keys_to_delete = []
    for line in lines:
        keys_to_delete.append(line)
        keys_to_delete.append('{}:executions'.format(line))
    r = Redis()
    for key in keys_to_delete:
        r.delete(key)


if __name__ == '__main__':
    if len(sys.argv) == 1 or sys.argv[1] == 'check':
        print('checking orphan tasks')
        check()
    elif sys.argv[1] == 'remove':
        if len(sys.argv) < 3:
            print('key file path is required')
            exit(-1)
        remove(sys.argv[2])
    else:
        print('not support')