heartbeat.py

import argparse
import json
import signal
import time
import timeit

start_time = timeit.default_timer()

parser = argparse.ArgumentParser()
parser.add_argument('--no-heartbeat', action='store_true', help='disable heartbeat logs')
parser.add_argument('--duration', type=int, help='runtime of script in seconds', required=True)
parser.add_argument('--ignore-signals', action='store_true', help='do not exit when signal is received')
parser.add_argument('--exit-code', type=int, default=0, help='exit code to return when the script finishes')
args = parser.parse_args()


def elapsed_time() -> str:
    return timeit.default_timer() - start_time


def log(event_type: str, data: dict = None):
    print(json.dumps({
        'event': event_type,
        'time': elapsed_time(),
        **(data or dict()),
    }))


def signal_handler(signum, frame):
    log('signal', {'signal_type': signal.Signals(signum).name})
    if not args.ignore_signals:
        log('stop')
        exit()


for sig in (signal.SIGINT, signal.SIGABRT, signal.SIGSEGV, signal.SIGTERM):
    signal.signal(sig, signal_handler)


for i in range(args.duration):
    if not args.no_heartbeat:
        log('heartbeat')
    time.sleep(1 + i + start_time - timeit.default_timer())

log('stop', {'exit_code': args.exit_code})
exit(args.exit_code)