#!/usr/bin/env python
"""
Bootstrax: XENONnT online processing manager
=============================================
How to use
----------------
    <activate conda environment>
    bootstrax --production
----------------

First draft: Jelle Aalbers, 2018
With additional input from: Joran Angevaare, 2020

This script watches for new runs to appear from the DAQ, then starts a strax process to process them. If a run fails, it will retry it with exponential backoff.

You can run more than one bootstrax instance, but only one per machine. If you start a second one on the same machine, it will try to kill the first one.


Philosophy
----------------
Bootstrax has a crash-only / recovery first philosophy. Any error in the core code causes a crash; there is no nice exit or mandatory cleanup. Bootstrax focuses on recovery after restarts: before starting work, we look for and fix any mess left by crashes.

This ensures that hangs and hard crashes do not require expert tinkering to repair databases. Plus, you can just stop the program with ctrl-c (or, in principle, pulling the machine's power plug) at any time.

Errors during run processing are assumed to be retry-able. We track the number of failures per run to decide how long to wait until we retry; only if a user marks a run as 'abandoned' (using an external system, e.g. the website) do we stop retrying.


Mongo documents
----------------
Bootstrax records its status in a document in the 'bootstrax' collection in the runs db. These documents contain:
  - **host**: socket.getfqdn()
  - **time**: last time this bootstrax showed life signs
  - **state**: one of the following:
    - **busy**: doing something
    - **idle**: NOT doing something; available for processing new runs

Additionally, bootstrax tracks information with each run in the 'bootstrax' field of the run doc. We could also put this elsewhere, but it seemed convenient. This field contains the following subfields:
  - **state**: one of the following:
    - **considering**: a bootstrax is deciding what to do with it
    - **busy**: a strax process is working on it
    - **failed**: something is wrong, but we will retry after some amount of time
    - **abandoned**: bootstrax will ignore this run
  - **reason**: reason for last failure, if there ever was one (otherwise this field does not exists). Thus, it's quite possible for this field to exist (and show an exception) when the state is 'done': that just means it failed at least once but succeeded later. Tracking failure history is primarily the DAQ log's reponsibility; this message is only provided for convenience.
   - **n_failures**: number of failures on this run, if there ever was one (otherwise this field does not exist).
  - **next_retry**: time after which bootstrax might retry processing this run. Like 'reason', this will refer to the last failure.
Finally, bootstrax outputs the load on the eventbuilder machine(s) whereon it is running to a collection in the DAQ database into the capped collection 'eb_monitor'. This collection contains information on what bootstrax is thinking of at the moment
  - **disk_used**: used part of the disk whereto this bootstrax instance is writing to (in percent)
"""
__version__ = '0.5.9'

import argparse
from datetime import datetime, timedelta, timezone
import logging
import multiprocessing
import npshmex
import os
import os.path as osp
import signal
import socket
import shutil
import time
import traceback
from tqdm import tqdm

import numpy as np
import pymongo
from psutil import pid_exists, disk_usage
import pytz
import strax
import straxen
import threading
import utilix

logging.basicConfig(level=logging.INFO,
                    format='%(relativeCreated)6d %(threadName)s %(name)s %(message)s')

parser = argparse.ArgumentParser(
    description="XENONnT online processing manager")
parser.add_argument('--debug', action='store_true',
                    help="Start strax processes with debug logging.")
parser.add_argument('--profile', type=str, default='false',
                    help="Option to run strax in profiling mode. "
                         "argument specifies the name of the profile if not 'false'. Use e.g. 'date'.prof")
parser.add_argument('--cores', type=int, default=8,
                    help="Maximum number of workers to use in a strax process. "
                         "Set to -1 for all available cores")
parser.add_argument('--target', default='event_info',
                    help="Strax data type name that should be produced")
parser.add_argument('--fix_target', action='store_true',
                    help="Don't allow bootstrax to switch to a different target for special runs")
parser.add_argument('--infer_mode', action='store_true',
                    help="Determine best number max-messages and cores for each run automatically. "
                         "Overrides --cores and --max_messages")
parser.add_argument('--delete_live', action='store_true',
                    help="Delete live_data after successful processing of the run.")
parser.add_argument('--production', action='store_true',
                    help="Run bootstrax in production mode. Assuming test mode otherwise to prevent "
                         "interactions with the runs-database")
parser.add_argument('--undying', action='store_true',
                    help="Except any error and ignore it")
parser.add_argument('--sub_d_targets', nargs='*',
                    default=['raw_records_he', 'raw_records_nv', 'raw_records_mv'],
                    help="Target(s) for other sub-detectors. If not produced automatically "
                         "when processing tpc data, st.make the requested data later.")
parser.add_argument('--max_messages', type=int,
                    default=10,
                    help="number of max mailbox messages")

actions = parser.add_mutually_exclusive_group()
actions.add_argument('--process', type=int, metavar='NUMBER',
                     help="Process a single run, regardless of its status.")
actions.add_argument('--fail', nargs='+',
                     metavar=('NUMBER', 'REASON'),
                     help="Fail run number, optionally with reason")
actions.add_argument('--abandon', nargs='+',
                     metavar=('NUMBER', 'REASON'),
                     help="Abandon run number, optionally with reason")

args = parser.parse_args()

##
# Configuration
##

print(f'---\n bootstrax version {__version__}\n---')

# The folder that can be used for testing bootstrax (i.e. non production mode). It will be written to:
test_data_folder = ('/nfs/scratch/bootstrax/' if
                    os.path.exists('/nfs/scratch/bootstrax/')
                    else './bootstrax/')

# Timeouts in seconds
timeouts = {
    # Waiting between escalating SIGTERM -> SIGKILL -> crashing bootstrax
    # when trying to kill another process (usually child strax)
    'signal_escalate': 3,
    # Minimum waiting time to retry a failed run
    # Escalates exponentially on repeated failures: 1x, 5x, 25x, 125x, 125x, 125x, ...
    # Some jitter is applied: actual delays will randomly be 0.5 - 1.5x as long
    'retry_run': 60,
    # Maximum time for strax to complete a processing
    # if exceeded, strax will be killed by bootstrax
    'max_processing_time': 7200,
    # Sleep between checking whether a strax process is alive
    'check_on_strax': 10,
    # Maximum time a run is 'busy' without a further update from
    # its responsible bootstrax. Bootstrax normally updates every
    # check_on_strax seconds, so make sure this is substantially
    # larger than check_on_strax.
    'max_busy_time': 120,
    # Maximum time a run is in the 'considering' state
    # if exceeded, will be labeled as an untracked failure
    'max_considering_time': 60,
    # Minimum time to wait between database cleanup operations
    'cleanup_spacing': 60,
    # Sleep time when there is nothing to do
    'idle_nap': 10,
    # If we don't hear from a bootstrax on another host for this long,
    # remove its entry from the bootstrax status collection
    # Must be much longer than idle_nap and check_on_strax!
    'bootstrax_presumed_dead': 300,
    # Ebs3-5 normally do all the processing. However if all are busy
    # for a longer period of time, the ebs0-2 can also help with
    # processing.
    'eb3-5_max_busy_time': 5 * 60,
    # Bootstrax writes it's state to the daq-database. To have a backlog we store this
    # state using a TTL collection. To prevent too many entries in this backlog, only
    # create new entries if the previous entry is at least this old (in seconds).
    'min_status_interval': 60
}

# The disk that the eb is writing to may fill up at some point. The data should
# be written to datamanager at some point. This may clean up data on the disk,
# hence, we can check if there is sufficient diskspace and if not, wait a while.
# Below are the max number of times and number of seconds bootstrax will wait.
wait_diskspace_max_space_percent = 90
wait_diskspace_n_max = 60 * 24 * 7  # times
wait_diskspace_dt = 60  # seconds
assert timeouts['bootstrax_presumed_dead'] > wait_diskspace_dt, "wait_diskspace_dt too large"

# Fields in the run docs that bootstrax uses
bootstrax_projection = f"name start end number bootstrax status mode " \
                       f"data.host data.type data.location " \
                       f"daq_config.processing_threads daq_config.compressor " \
                       f"daq_config.strax_fragment_payload_bytes " \
                       f"daq_config.strax_chunk_length daq_config.strax_chunk_overlap".split()

# Filename for temporary storage of the exception
# This is used to communicate the exception from the strax child process
# to the bootstrax main process
exception_tempfile = 'last_bootstrax_exception.txt'

# The name of the thread that is opened to delete live_data
delete_thread_name = 'DeleteThread'

# boostrax state for 'dead' or old entries in the bs_coll
dead_state = 'dead_bootstrax'

# The maximum time difference (s) allowed between the timestamps in the data and the
# duration of the run (from the runs metadeta). Fail if the difference is larger than:
max_timetamp_diff = 5

# The maximum number of retries for processing a run. After this many times of retrying
# to process a run, the DAQ-group has to either manually fix this run or manually fail it.
max_n_retry = 20

# Bootstrax retries runs multiple times. If there have been this many fails we could lower
# the resources to be somewhat more lenient on the CPU and RAM of the eventbuilders. Use
# this option with care as we might be running in a sub-optimal mode that may not be
# noticed and eventbuilders may be spending more time on trying to reprocess failed runs.
lower_resources_after_n_failures = 10


##
# Initialize globals (e.g. rundb connection)
##

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(name)s %(levelname)-8s %(message)s',
    datefmt='%m-%d %H:%M')
log = logging.getLogger()
hostname = socket.getfqdn()

# Set the output folder
output_folder = '/data/xenonnt_processed/'

if not args.production:
    # This means we are in some test mode
    wait_diskspace_max_space_percent = 80
    output_folder = test_data_folder
    if not os.path.exists(output_folder):
        log.warning(f'Creating {output_folder}')
        os.mkdir(output_folder)

    log.warning(
        f'\n---------------'
        f'\nBe aware, bootstrax not running in production mode. Specify with --production.'
        f'\nWriting new data to {output_folder}. Not saving this location in the runsDB.'
        f'\nNot writing to the runs-database.'
        f'\n---------------')
    time.sleep(5)
else:
    if not args.delete_live:
        log.warning("Production mode is designed to run with '--delete_live'\nplease restart bootstrax")
    if not args.infer_mode:
        log.warning("Better performance is expected in production mode with '--infer_mode'\nplease restart bootstrax")


def new_context(cores=args.cores, max_messages=args.max_messages, timeout=300):
    """Create strax context that can access the runs db"""
    # We use exactly the same logic of straxen to access the runs DB;
    # this avoids duplication, and ensures strax can access the runs DB if we can
    context = straxen.contexts.xenonnt_online(
        output_folder=output_folder,
        we_are_the_daq=True,
        allow_multiprocess=cores > 1,
        allow_shm=cores > 1,
        max_messages=max_messages,
        timeout=timeout)
    if not args.production:
        context.storage = [strax.DataDirectory(output_folder)]
    return context


st = new_context()

# DAQ database
daq_db_name = 'daq'
daq_uri = straxen.get_mongo_uri(header='rundb_admin',
                                user_key='mongo_daq_username',
                                pwd_key='mongo_daq_password',
                                url_key='mongo_daq_url')
daq_client = pymongo.MongoClient(daq_uri)
daq_db = daq_client[daq_db_name]
bs_coll = daq_db['eb_monitor']
ag_stat_coll = daq_db['aggregate_status']
log_coll = daq_db['log']

# Runs database
run_dbname = straxen.uconfig.get('rundb_admin', 'mongo_rdb_database')
run_collname = 'runs'
if args.production:
    run_db = st.storage[0].client[run_dbname]
else:
    # Please note, this is a read only account on the rundb
    run_uri = straxen.get_mongo_uri()
    run_client = pymongo.MongoClient(run_uri)
    run_db = run_client[run_dbname]
run_coll = run_db[run_collname]
run_db.command('ping')

run_coll = run_db[run_collname]

# Ping the databases to ensure the mongo connections are working
if not args.undying:
    run_db.command('ping')
    daq_db.command('ping')


def main():
    # Check that writing access is OK, otherwise report to the database and die
    if os.access(output_folder, os.W_OK) is not True:
        message = f'No writing access to {output_folder}'
        log_warning(message, priority='fatal')
        raise IOError(message)

    if args.cores == -1:
        # Use all of the available cores on this machine
        args.cores = multiprocessing.cpu_count()
        log.info(f'Set cores to n_tot, using {args.cores} cores')

    if args.fail:
        args.fail += ['']  # Provide empty reason if none specified
        manual_fail(number=int(args.fail[0]), reason=args.fail[1])

    elif args.abandon:
        number = int(args.abandon[0])
        if len(args.abandon) > 1:
            manual_fail(number=number, reason=args.abandon[1])
        abandon(number=number)

    elif args.process:
        t_start = now()
        number = args.process
        rd = consider_run({'number': number})
        if rd is None:
            message = f"Trying to process single run but no run numbered {number} exists"
            log_warning(message, priority='fatal')
            raise ValueError(message)
        process_run(rd)
        log.info(f'bootstrax ({hostname}) finished run {number} in {(now() - t_start).seconds} seconds')
        wait_on_delete_thread()

    else:
        # Start processing
        main_loop()


##
# Main loop
##

def main_loop():
    """Infinite loop looking for runs to process"""

    # Wait for mongo connection
    ping_dbs()

    # Ensure we're the only bootstrax on this host
    any_other_running = list(bs_coll.find({'host': hostname,
                                           'pid': {'$ne': os.getpid()}}))
    for x in any_other_running:
        if pid_exists(x['pid']) and x['pid']:
            log.warning(f'Bootstrax already running with PID {x["pid"]}, trying to kill it.')
            kill_process(x['pid'])

    # # Register ourselves
    set_state('starting')
    t_start = now()

    next_cleanup_time = now()
    # keep track of the ith run that we have seen when we are not in production mode
    new_runs_seen, failed_runs_seen = 0, 1
    while True:
        log.info(f'bootstrax running for {(now() - t_start).seconds} seconds')
        sufficient_diskspace()
        log.info("Looking for work")
        if not eb_can_process():
            # Nothing to do, let's do some cleanup and set to idle
            delete_temp_files()
            set_state('idle')
            time.sleep(timeouts['idle_nap'])
            continue
        set_state('busy')
        # Check resources are still OK, otherwise crash / reboot program

        # Process new runs
        rd = consider_run({"bootstrax.state": {"$exists": False}},
                          test_counter=new_runs_seen)
        if rd is not None:
            new_runs_seen += 1
            process_run(rd)
            continue

        # Scan DB for runs with unusual problems
        if now() > next_cleanup_time:
            cleanup_db()
            next_cleanup_time = now(plus=timeouts['cleanup_spacing'])

        # Any failed runs to retry?
        # Only try one run, we want to be back for new runs quickly
        rd = consider_run({"bootstrax.state": 'failed',
                           "bootstrax.n_failures": {'$lt': max_n_retry},
                           "bootstrax.next_retry": {'$lt': now()}
                           }, test_counter=failed_runs_seen)
        if rd is not None:
            failed_runs_seen += 1
            process_run(rd)
            continue
        # Nothing to do, let's do some cleanup
        delete_temp_files()
        if not args.production:
            log.info(f'We have gone through the rundDB in a readonly mode there are no '
                     f'runs left. We looked at {new_runs_seen} new runs and '
                     f'{failed_runs_seen} previously failed runs.')
            break
        log.info("No work to do, waiting for new runs or retry timers")
        set_state('idle')
        time.sleep(timeouts['idle_nap'])


##
# General helpers
##


def now(plus=0):
    return datetime.now(pytz.utc) + timedelta(seconds=plus)


def kill_process(pid, wait_time=None):
    """Kill process pid, or raise RuntimeError if we cannot
    :param wait_time: time to wait before escalating signal strength
    """
    if wait_time is None:
        wait_time = timeouts['signal_escalate']
    if not pid_exists(pid):
        return

    for sig in [signal.SIGTERM, signal.SIGKILL, 'die']:
        time.sleep(wait_time)
        if not pid_exists(pid):
            return
        if signal == 'die':
            message = f"Could not kill process {pid}"
            log_warning(message, priority='fatal')
            raise RuntimeError(message)
        os.kill(pid, sig)


def override_target(rd):
    """
    Check if the target should be overridden based on the mode of the DAQ for this run
    :param rd: rundoc
    :return: override": False or LED or raw_records}
    """
    if args.fix_target:
        return False

    # Special modes override target for these
    led_modes = ['pmtgain']
    diagnostic_modes = ['exttrig', 'noise', 'pmtap']

    mode = str(rd.get('mode'))
    log.info(f'override_target::\tmode is {mode}, changing target if needed')
    if np.any([m in mode for m in led_modes]):
        return 'led_calibration'
    elif np.any([m in mode for m in diagnostic_modes]):
        return 'raw_records'
    else:
        return False


def set_state(state, update_fields=None):
    """Inform the bootstrax collection we're in a different state

    if state is None, leave state unchanged, just update heartbeat time
    """
    ping_dbs()

    # Find the last message of this host
    previous_entry = bs_coll.find_one({'host': hostname},
                                      sort=[('time', pymongo.DESCENDING)])
    if not state:
        state = 'None' if not previous_entry else previous_entry.get('state')

    bootstrax_state = dict(
        host=hostname,
        pid=os.getpid(),
        time=now(),
        state=state,
        target=args.target,
        max_cores=args.cores,
        max_messages=args.max_messages,
        undying=args.undying,
        production_mode=args.production
    )
    if update_fields:
        bootstrax_state.update(update_fields)
    if (not previous_entry or
            (now() - previous_entry['_id'].generation_time).seconds > timeouts['min_status_interval']):
        bs_coll.insert_one(bootstrax_state)
    else:
        bs_coll.update_one({'_id': previous_entry['_id']}, {'$set': bootstrax_state})


def send_heartbeat(update_fields=None):
    """Inform the bootstrax collection we're still here
    Use during long-running tasks where state doesn't change
    """
    # Same as set_state, just don't change state
    set_state(None, update_fields=update_fields)


def log_warning(message, priority='warning', run_id=None):
    """Report a warning to the terminal (using the logging module)
    and the DAQ log DB.
    :param message: insert string into log_coll
    :param priority: severity of warning. Can be:
        info: 1,
        warning: 2,
        <any other valid python logging level, e.g. error or fatal>: 3
    :param run_id: optional run id.
    """
    ping_dbs()
    if not args.production:
        return
    getattr(log, priority)(message)
    # Log according to redax rules
    # https://github.com/coderdj/redax/blob/master/MongoLog.hh#L22
    warning_message = {
        'message': message,
        'user': f'bootstrax_{hostname}',
        'priority':
            dict(debug=0,
                 info=1,
                 warning=2,
                 error=3,
                 fatal=4,
                 ).get(priority.lower(), 3)}
    if run_id is not None:
        warning_message.update({'runid': int(run_id)})
    log_coll.insert_one(warning_message)


def eb_can_process():
    """"The new ebs (eb3-5) should be sufficient to process all data. In exceptional
    circumstances eb3-5 cannot keep up. Only let eb0-2 also process data in such cases.
    Before eb0-2 are also used for processing two criteria have to be fulfilled:
      - There should be runs waiting to be processed
      - Eb3-5 should be busy processing for a substantial time.
    :returns: bool if this host should process a run"""

    # Check mongo connection
    ping_dbs()

    # eb3-5 always process.
    if hostname in ['eb3.xenon.local', 'eb4.xenon.local', 'eb5.xenon.local']:
        return True

    # In test mode we can always process
    if not args.production:
        return True
    elif 'eb2' in hostname:
        log_warning('Why is eb2 alive?!', priority='error')
        return False

    # Check that there are runs that are waiting to be processed. If there are few, this
    # eb should not process.
    max_queue_new_runs = 2

    # Count number of runs untouched by bootstrax.
    n_untouched_runs = run_coll.count_documents({'bootstrax': {'$exists': 0}})

    # Check that eb3-5 are all busy for at least some time.
    n_ebs_running = 0
    n_ebs_busy = 0
    for eb_i in range(3, 6):
        # Should count if eb3-5 are registered as running (as one might be offline).
        bootstrax_on_host = bs_coll.find_one({
            'host': f'eb{eb_i}.xenon.local',
            'time': {'$gt': now(-timeouts['bootstrax_presumed_dead'])}},
            sort=[('time', pymongo.DESCENDING)])

        if bootstrax_on_host:
            n_ebs_running += 1
            running_eb = run_coll.find_one({
                    'bootstrax.state': 'busy',
                    'bootstrax.host': f'eb{eb_i}.xenon.local',
                    'bootstrax.started_processing': {
                        '$lt': now(-timeouts['eb3-5_max_busy_time'])}})
            if running_eb:
                n_ebs_busy += 1
                log.info(f'eb_can_process::\t eb{eb_i} is busy')
    if ((n_ebs_running == n_ebs_busy and n_untouched_runs > max_queue_new_runs) or
            not n_ebs_running):
        # There is a need for this eb to also process data (for now).
        log.info(f'eb_can_process::\tThere is a need for {hostname} to process data. As we'
                 f' have {n_ebs_running} running and {n_ebs_busy} busy')
        return True
    elif n_untouched_runs and n_untouched_runs < max_queue_new_runs:
        log.info(f'eb_can_process::\tDo not process on {hostname} as there are few new runs.')
    else:
        log.info(f'eb_can_process::\tDo not process on {hostname}, new ebs are taking care')
    log.info(f'eb_can_process::\trunning: {n_ebs_running}\tbusy: {n_ebs_busy}\tqueue: {n_untouched_runs}')
    return False


def infer_mode(rd):
    """Infer a safe operating mode of running bootstrax based on the size of the first
    chunk. Estimating save parameters for running bootstrax from:
    https://xe1t-wiki.lngs.infn.it/doku.php?id=xenon:xenonnt:dsg:daq:eb_speed_tests_update
    returns: dictionary of how many cores and max_messages should be used based on an
    estimated data rate.
    """
    default_opt = dict(cores=args.cores, max_messages=args.max_messages, timeout=300)

    # Get data rate from dispatcher
    try:
        docs = ag_stat_coll.aggregate([
                    {'$match': {'number': rd['number']}},
                    {'$group': {'_id': '$detector', 'rate': {'$max': '$rate'}}}
                ])
        data_rate = sum([d['rate'] for d in docs])
    except Exception as e:
        log_warning(f'infer_mode ran into {e}. Cannot infer mode, using default mode.',
                    run_id=f'{rd["number"]:06}', priority='info')
        data_rate = None

    # Find out if eb is new (eb3-eb5):
    is_new_eb = int(hostname[2]) >= 3  # ebX.xenon.local
    log.info(f"infer_mode::\teb{int(hostname[2])}")
    eb = 'new_eb' if is_new_eb else 'old_eb'

    # Return a run-mode that is empirically found to provide stable but fast processing.
    log.info(f'infer_mode::\tWorking with a data rate of {data_rate:.1f} MB/s')
    # TODO
    #  decrease timeout later. Unsure why needed now.
    run_settings = {
        'new_eb':
            {'low_rate': dict(cores=24, max_messages=20, timeout=200),
             'med_rate': dict(cores=24, max_messages=20, timeout=400),
             'high_rate': dict(cores=20, max_messages=15, timeout=1000),
             'very_high_rate': dict(cores=20, max_messages=15, timeout=1500),
             'max_rate': dict(cores=16, max_messages=10, timeout=2000)},
        'old_eb':
            {'low_rate': dict(cores=30, max_messages=20, timeout=200),
             'med_rate': dict(cores=20, max_messages=15, timeout=400),
             'high_rate': dict(cores=10, max_messages=15, timeout=1000),
             'very_high_rate': dict(cores=10, max_messages=10, timeout=1500),
             'max_rate': dict(cores=8, max_messages=10, timeout=2000)}}
    if data_rate is None:
        result = default_opt
    elif data_rate < 100:
        result = run_settings[eb]['low_rate']
    elif data_rate < 200:
        result = run_settings[eb]['med_rate']
    elif data_rate < 500:
        result = run_settings[eb]['high_rate']
    elif data_rate < 700:
        result = run_settings[eb]['very_high_rate']
    else:
        result = run_settings[eb]['max_rate']
    log.info(f'infer_mode::\tOverride processing mode for run {rd["number"]} changing to:'
             f'\t {result}')

    # Lower the settings if there are continuous failures on this run:
    if rd.get('bootstrax', {}).get('n_failures', 0) >= lower_resources_after_n_failures:
        result = dict(cores=max(4, result['cores'] // 2),
                      max_messages=max(4, result['max_messages'] // 2),
                      timeout=min(1000, result['timeout'] * 1.5)
                      )
        log_warning(f'infer_mode::\tRepeated failures. Lowering mode to {result}. This '
                    f'may indicated a sub-optimal state of {hostname}!',
                    priority='info',
                    run_id=f'{rd["number"]:06}')

    return result


##
# Host interactions
##

def sufficient_diskspace():
    """Check if there is sufficient space available on the local disk to write to"""
    for i in range(wait_diskspace_n_max):
        disk_pct = disk_usage(output_folder).percent
        if disk_pct < wait_diskspace_max_space_percent:
            log.info(f'Check disk space: {disk_pct:.1f}% full')
            # Sufficient space to write to, let's continue
            return
        else:
            log.info(f'Insufficient free disk space ({disk_pct:.1f}% full) '
                     f'on {hostname}. Waiting {wait_diskspace_dt} s ({i}th iteration)')
            time.sleep(wait_diskspace_dt)
            send_heartbeat(dict(state='disk full'))
    set_state(dead_state)
    message = f"No disk space to write to. Kill bootstrax on {hostname}"
    log_warning(message, priority='fatal')
    raise RuntimeError(message)


def delete_temp_files():
    """
    removes _temp files from output_folder if successfully processed
    """
    # Check mongo connection
    ping_dbs()

    temp_files = [f for f in os.listdir(output_folder) if f.endswith('_temp')]
    for f in temp_files:
        run_id = int(f.split('-')[0])
        if run_coll.find_one(
                {'bootstrax.state': 'done',
                 'number': run_id},
                {'_id': 1}):
            log.info(f'removing {output_folder}/{f}')
            shutil.rmtree(f'{output_folder}/{f}')


def delete_live_data(rd, live_data_path):
    """
    Open thread to delete the live_data
    """
    if args.production and os.path.exists(live_data_path) and args.delete_live:
        delete_thread = threading.Thread(name=delete_thread_name,
                                         target=_delete_data,
                                         args=(rd, live_data_path, 'live'))
        log.info(f'Starting thread to delete {live_data_path} at {now()}')
        # We rather not stop deleting the live_data if something else fails. Set the thread to daemon.
        delete_thread.setDaemon(True)
        delete_thread.start()
        log.info(f'DeleteThread {live_data_path} should be running in parallel, '
                 f'continue MainThread now: {now()}')


def _delete_data(rd, path, data_type):
    """After completing the processing and updating the runsDB, remove the
    live_data"""
    ping_dbs()

    if data_type == 'live' and not args.delete_live and args.production:
        message = 'Unsafe operation. Trying to delete live data!'
        log_warning(message, priority='fatal')
        raise ValueError(message)
    log.info(f'Deleting data at {path}')
    if os.path.exists(path):
        shutil.rmtree(path)
    log.info(f'deleting {path} finished')
    # Remove the data location from the rundoc and append it to the 'deleted_data' entries
    if not os.path.exists(path):
        log.info('changing data field in rundoc')
        for ddoc in rd['data']:
            if ddoc['type'] == data_type:
                break
        for k in ddoc.copy().keys():
            if k in ['location', 'meta', 'protocol']:
                ddoc.pop(k)

        ddoc.update({'at': now(), 'by': hostname})
        log.info(f'update with {ddoc}')
        run_coll.update_one({'_id': rd['_id']},
                            {"$addToSet": {'deleted_data': ddoc},
                             "$pull": {"data":
                                           {"type": data_type,
                                            "host": {'$in': ['daq', hostname]}}}})
    else:
        message = f"Something went wrong we wanted to delete {path}!"
        log_warning(message, priority='fatal')
        raise ValueError(message)


def wait_on_delete_thread():
    """
    Check that the thread with the delete_thread_name is finished before continuing.
    """
    threads = threading.enumerate()
    for thread in threads:
        if thread.name == delete_thread_name:
            wait = True
            while wait:
                wait = False
                if thread.isAlive():
                    log.info(f'{thread.name} still running take a {timeouts["idle_nap"]} s nap')
                    time.sleep(timeouts['idle_nap'])
                    wait = True
    log.info(f'Checked that {delete_thread_name} finished')


def clear_shm():
    """Manually delete files in /dev/shm/ created by npshmex on starup."""
    shm_dir = '/dev/shm/'
    shm_files = [f for f in os.listdir(shm_dir) if 'npshmex' in f]

    if not len(shm_files):
        return
    log.info(f'clear_shm:: clearing {len(shm_files)} files')
    for f in tqdm(shm_files):
        os.remove(shm_dir + f)


##
# Run DB interaction
##


def ping_dbs():
    while True:
        try:
            run_db.command('ping')
            daq_db.command('ping')
            break
        except Exception as ping_error:
            log.warning(f'Failed to connect to Mongo. Ran into {ping_error}. Sleep for a '
                        f'minute.')
            time.sleep(60)


def get_run(*, mongo_id=None, number=None, full_doc=False):
    """Find and return run doc matching mongo_id or number
    The bootstrax state is left unchanged.

    :param full_doc: If true (default is False), return the full run doc
        rather than just fields used by bootstrax.
    """
    ping_dbs()
    if number is not None:
        query = {'number': number}
    elif mongo_id is not None:
        query = {'_id': mongo_id}
    else:
        # This means you are not running a normal bootstrax (no reason to report to rundb)
        raise ValueError("Please give mongo_id or number")

    return run_coll.find_one(query, projection=None if full_doc else bootstrax_projection)


def set_run_state(rd, state, return_new_doc=True, **kwargs):
    """Set state of run doc rd to state
    return_new_doc: if True (default), returns new document.
        if False, instead returns the original (un-updated) doc.

    Any additional kwargs will be added to the bootstrax field.
    """
    ping_dbs()
    if not args.production:
        return run_coll.find_one({'_id': rd['_id']})

    bd = rd['bootstrax']
    bd.update({
        'state': state,
        'host': hostname,
        'time': now(),
        **kwargs})

    if state == 'failed':
        bd['n_failures'] = bd.get('n_failures', 0) + 1

    return run_coll.find_one_and_update(
        {'_id': rd['_id']},
        {'$set': {'bootstrax': bd}},
        return_document=return_new_doc,
        projection=bootstrax_projection)


def check_data_written(rd):
    """
    checks that the data as written in the runs-database is actually
    available on this machine
    :param rd: rundoc
    :return: type bool, False if not all paths exist or if there are no files
    on this host
    """
    files_written = 0
    for ddoc in rd['data']:
        if ddoc['host'] == hostname:
            if os.path.exists(ddoc['location']):
                files_written += 1
            else:
                return False
    return files_written > 0


def all_files_saved(rd, wait_max=600, wait_per_cycle=10):
    """
    Check that all files are written. It might be that the savers are still in
    the process of renaming from folder_temp to folder. Hence allow some wait
    time to allow the savers to finish
    :param rd: rundoc
    :param wait_max: max seconds to wait for data to save
    :param wait_per_cycle: wait this many seconds if the data is not yet there
    """
    start = time.time()
    while not check_data_written(rd):
        if time.time() - start > wait_max:
            return False
        send_heartbeat()
        time.sleep(wait_per_cycle)
    return True


def upload_file_metadata(rd):
    """
    Update the data-field in the rundoc with a portion of the metadata. Also count the
    number of files on the location on the basis of the data entry in the rundoc. The
    filecount info is used for Admix to checksum that all the files are correctly uploaded
    to Rucio.
    :param rd: rundoc
    """
    try:
        st_meta = new_context(cores=args.cores, max_messages=args.max_messages, timeout=100)
        st_meta.set_context_config({'forbid_creation_of': '*'})
    except Exception as e:
        log_warning(f"Cannot create context to read the metadata: {e}", priority='warning')
        st_meta = None

    for ddoc in rd['data']:
        if hostname != ddoc['host']:
            continue
        loc = ddoc.get('location', '')

        if os.path.exists(loc):
            file_count = len(os.listdir(loc))
            # Can also get the latter from st.lineage_for but too lazy for that
            data_size_mb, avg_data_size_mb, lineage_hash = None, None, None
            run_id = '%06d' % rd['number']
            if st_meta is not None:
                try:
                    md = st.get_meta(run_id, ddoc['type'])
                    chunk_mb = [chunk['nbytes']/(1e6) for chunk in md['chunks']]
                    data_size_mb = int(np.sum(chunk_mb))
                    avg_data_size_mb = int(np.average(chunk_mb))
                    lineage_hash = md['lineage_hash']
                except Exception as e:
                    log_warning(f"Cannot load metadata of {ddoc['type']}: {e}",
                                priority='warning',
                                run_id=f'{rd["number"]:06}')

            run_coll.update_one(
                    {'_id': rd['_id'],
                     'data.location': ddoc['location']},
                    {'$set':
                         {'data.$.file_count': file_count,
                          'data.$.meta.strax_version': strax.__version__,
                          'data.$.meta.straxen_version': straxen.__version__,
                          'data.$.meta.size_mb': data_size_mb,
                          'data.$.meta.avg_chunk_mb': avg_data_size_mb,
                          'data.$.meta.lineage_hash': lineage_hash}
                     })


def set_status_finished(rd):
    """Set the status to ready to upload for datamanager and admix"""
    # Check mongo connection
    ping_dbs()

    if not args.production:
        # Don't update the status if we are not in production mode
        return

    # First check that all the data is available (that e.g. no _temp files
    # are being renamed). This line should be over-redundant as we already
    # check earlier.
    all_files_saved(rd)

    # Only update the status if it does not exist or if it needs to be uploaded
    ready_to_upload = {'status': 'eb_ready_to_upload'}
    if rd.get('status') in [None, 'needs_upload']:
        run_coll.update_one(
            {'_id': rd['_id']},
            {'$set': ready_to_upload})
    elif rd.get('status') == ready_to_upload.get('status'):
        # This is strange, bootstrax already finished this run before
        log_warning('WARNING: bootstax has already marked this run as ready '
                    'for upload. Doing nothing.',
                    priority='warning',
                    run_id=f'{rd["number"]:06}')
    else:
        # Do not override this field for runs already uploaded in admix
        message = (f'Trying to set set the status {rd.get("status")} to '
                   f'{ready_to_upload}! One should not override this field.')
        log_warning(message, priority='fatal')
        raise ValueError(message)


def abandon(*, mongo_id=None, number=None):
    """Mark a run as abandoned"""
    set_run_state(
        get_run(mongo_id=mongo_id, number=number),
        'abandoned')


def consider_run(query, return_new_doc=True, test_counter=0):
    """Return one run doc matching query, and simultaneously set its bootstraxstate to 'considering'"""
    # Check mongo connection
    ping_dbs()

    # We must first do an atomic find-and-update to set the run's state
    # to "considering", to ensure the run doesn't get picked up by a
    # bootstrax on another host.
    if args.production:
        rd = run_coll.find_one_and_update(
            query,
            {"$set": {'bootstrax.state': 'considering'}},
            projection=bootstrax_projection,
            return_document=True,
            sort=[('start', pymongo.DESCENDING)])
        # Next, we can update the bootstrax entry properly with set_run_state
        # (adding hostname, time, etc.)
        if rd is None:
            return None
        return set_run_state(rd, 'considering', return_new_doc=return_new_doc)
    else:
        # Don't change the runs-database for test modes
        try:
            rds = run_coll.find(
                query,
                projection=bootstrax_projection,
                sort=[('start', pymongo.DESCENDING)])
            return rds[test_counter]
        except IndexError:
            return None


def fail_run(rd, reason, error_traceback=''):
    """Mark the run represented by run doc rd as failed with reason"""
    if 'number' not in rd:
        long_run_id = f"run <no run number!!?>:{rd['_id']}"
    else:
        long_run_id = f"run {rd['number']}:{rd['_id']}"

    # No bootstrax info is present when manually failing a run with args.fail
    if 'bootstrax' not in rd.keys():
        rd['bootstrax'] = {}
        rd['bootstrax']['n_failures'] = 0

    if 'n_failures' in rd['bootstrax'] and rd['bootstrax']['n_failures'] > 0:
        fail_name = 'Repeated failure'
        failure_message_level = 'info'
    else:
        fail_name = 'New failure'
        failure_message_level = 'warning'

    # Cleanup any data associated with the run
    # TODO: This should become optional, or just not happen at all,
    # after we're done testing (however, then we need some other
    # pruning mechanism)
    clean_run(mongo_id=rd['_id'])

    # Report to run db
    # It's best to do this after everything is done;
    # as it changes the run state back away from 'considering', so another
    # bootstrax could conceivably pick it up again.
    set_run_state(rd, 'failed',
                  reason=reason + error_traceback,
                  next_retry=(
                      now(plus=(timeouts['retry_run']
                                * np.random.uniform(0.5, 1.5)
                                # Exponential backoff with jitter
                                * 5 ** min(rd['bootstrax'].get('n_failures', 0), 3)
                                ))))

    # Report to DAQ log and screen. Let's not also add the entire traceback
    log_warning(f"{fail_name} on {long_run_id}: {reason}",
                priority=failure_message_level,
                run_id=f'{rd["number"]:06}')


def manual_fail(*, mongo_id=None, number=None, reason=''):
    """Manually mark a run as failed based on mongo_id or run number"""
    rd = get_run(mongo_id=mongo_id, number=number)
    fail_run(rd, "Manually set failed state. " + reason)


def get_compressor(rd, default_compressor="lz4"):
    """Read the compressor method from the run_doc. Return 'lz4' if no
    information is specified in the run_doc"""
    try:
        return rd["daq_config"]["compressor"]
    except KeyError:
        log_warning(f"Bootstrax couldn't read the compressor form the run_doc. "
                    f"Assuming 'lz4' for now",
                    priority='info',
                    run_id=f'{rd["number"]:06}')
        return default_compressor


##
# Processing
##

def run_strax(run_id, input_dir, target, n_readout_threads, compressor,
              run_start_time, samples_per_record, process_mode, target_override,
              daq_chunk_duration, daq_overlap_chunk_duration, n_fails,
              debug=False):
    # Check mongo connection
    ping_dbs()
    # Clear the swap memory used by npshmmex
    npshmex.shm_clear()
    # double check by forcefully clearing shm
    clear_shm()

    if debug:
        log.setLevel(logging.DEBUG)
    try:
        log.info(f"Starting strax to make {run_id} with input dir {input_dir}")

        # Never override the target for raw_records because we might be in a save mode
        if target != 'raw_records' and target_override:
            log.info(f'This is an non-standard run. Changing target from {target} to '
                     f'{target_override}.\nTo disable specify --fix_target')
            if target_override == 'led_calibration':
                # Increase the timeout a little
                process_mode['timeout'] = process_mode['timeout'] * 5
            target = target_override

        st = new_context(**process_mode)

        # Make a function for running strax, call the function to process the run
        # This way, it can also be run inside a wrapper to profile strax
        def st_make():
            """Run strax"""
            strax_config = dict(daq_input_dir=input_dir,
                                daq_compressor=compressor,
                                run_start_time=run_start_time,
                                record_length=samples_per_record,
                                daq_chunk_duration=daq_chunk_duration,
                                daq_overlap_chunk_duration=daq_overlap_chunk_duration,
                                n_readout_threads=n_readout_threads,
                                check_raw_record_overlaps=False,
                                )

            st.make(run_id, target,
                    config=strax_config,
                    max_workers=process_mode['cores'])

            if (target != 'raw_records' and n_fails == 0) or args.fix_target:
                # Make the nv, he, and mv data type only if we are:
                #  - Not processing up to raw_records (that means we are in some recovery mode)
                #  - We haven't failed this run before.
                # or
                #  - Are fixing the target (this is useful for testing but normally not used)
                for sub_d_target in args.sub_d_targets:
                    if sub_d_target not in st._plugin_class_registry:
                        log_warning(f'Trying to make unknown data type {sub_d_target}',
                                    priority='info',
                                    run_id=run_id)
                        continue
                    elif not st.is_stored(run_id, sub_d_target):
                        st.make(run_id, sub_d_target,
                                config=strax_config,
                                max_workers=process_mode['cores'])

        if args.profile.lower() == 'false':
            st_make()
        else:
            prof_file = f'run{run_id}_{args.profile}'
            if '.prof' not in prof_file:
                prof_file += '.prof'
            log.info(f'starting with profiler, saving as {prof_file}')
            with strax.profile_threaded(prof_file):
                st_make()
    except Exception as e:
        # Write exception to file, so bootstrax can read it
        exc_info = strax.formatted_exception()
        with open(exception_tempfile, mode='w') as f:
            f.write(exc_info)
        with open(f'./bootstrax_exceptions/{run_id}_exception.txt', mode='w') as f:
            f.write(exc_info)
        raise


def process_run(rd, send_heartbeats=args.production):
    # Check mongo connection
    ping_dbs()  # TODO obsolete?

    log.info(f"Starting processing of run {rd['number']}")
    if rd is None:
        raise RuntimeError("Pass a valid rundoc, not None!")

    # Shortcuts for failing
    class RunFailed(Exception):
        pass

    def fail(reason, **kwargs):
        if args.production:
            fail_run(rd, reason, **kwargs)
        else:
            log.warning(reason)
        raise RunFailed

    try:

        try:
            run_id = '%06d' % rd['number']
        except Exception as e:
            fail(f"Could not format run number: {str(e)}")

        for dd in rd['data']:
            if 'type' not in dd:
                fail("Corrupted data doc, found entry without 'type' field")
            if dd['type'] == 'live':
                break
            elif not args.production:
                # We are just testing let's assume its on the usual location
                dd = {'type': 'live', 'location': '/live_data/xenonnt/', 'host': 'daq'}
            else:
                fail("Non-live data already registered; untracked failure?")
        else:
            if not args.production:
                # We are just testing let's assume its on the usual location
                dd = {'type': 'live', 'location': '/live_data/xenonnt/', 'host': 'daq'}
            else:
                fail(f"No live data entry in rundoc")

        if not osp.exists(dd['location']):
            fail(f"No access to live data folder {dd['location']}")

        if 'daq_config' not in rd:
            fail('No daq_config in the rundoc!')
        try:
            # Fetch parameters from the rundoc. If not readable, let's use redax' default
            # values (that are hardcoded here).
            dq_conf = rd['daq_config']
            to_read = ('processing_threads', 'strax_chunk_length', 'strax_chunk_overlap',
                       'strax_fragment_payload_bytes')
            report_missing_config = [conf for conf in to_read if conf not in dq_conf]
            if report_missing_config:
                log_warning(f'{", ".join(report_missing_config)} not in rundoc for '
                            f'{run_id}! Using default values.',
                            priority='info',
                            run_id=run_id)
            thread_info = dq_conf.get('processing_threads', dict())
            n_readout_threads = sum([v for v in thread_info.values()])
            daq_chunk_duration = int(dq_conf.get('strax_chunk_length', 5) * 1e9)
            daq_overlap_chunk_duration = int(dq_conf.get('strax_chunk_overlap', 0.5) * 1e9)
            # note that value in rd in bytes hence //2
            samples_per_record = dq_conf.get('strax_fragment_payload_bytes', 220) // 2
            if not samples_per_record == 110:
                log.info(f'Samples_per_record = {samples_per_record}')
        except Exception as e:
            fail(f"Could not find {to_read} in rundoc: {str(e)}")

        if not n_readout_threads:
            fail(f"Run doc for {run_id} has no readout thread count info")

        loc = osp.join(dd['location'], run_id)
        if not osp.exists(loc):
            fail(f"No live data at claimed location {loc}")

        # Remove any previous processed data
        # If we do not do this, strax will just load this instead of
        # starting a new processing
        if args.production:
            clean_run(mongo_id=rd['_id'])
        else:
            clean_run_test_data(run_id)

        # Remove any temporary exception info from previous runs
        if osp.exists(exception_tempfile):
            os.remove(exception_tempfile)

        target = args.target
        n_fails = rd['bootstrax'].get('n_failures', 0)
        if not args.production and 'bootstrax' not in rd:
            # Bootstrax does not register in non-production mode
            pass
        elif n_fails > 1 and not args.process:
            # Failed before, and on autopilot: do just raw_records
            target = 'raw_records'

        compressor = get_compressor(rd)
        target_override = override_target(rd)
        try:
            run_start_time = rd['start'].replace(tzinfo=timezone.utc).timestamp()
        except Exception as e:
            fail(f"Could not find start in datetime.datetime object: {str(e)}")

        if args.infer_mode:
            process_mode = infer_mode(rd)
        else:
            process_mode = dict(cores=args.cores, max_messages=args.max_messages, timeout=300)

        strax_proc = multiprocessing.Process(
            target=run_strax,
            args=(run_id, loc, target, n_readout_threads, compressor,
                  run_start_time, samples_per_record, process_mode, target_override,
                  daq_chunk_duration, daq_overlap_chunk_duration, n_fails,
                  args.debug))

        t0 = now()
        info = dict(started_processing=t0)
        strax_proc.start()

        while True:
            if send_heartbeats:
                update = process_mode.copy()
                update.update(dict(run_id=run_id,
                                   target=target_override if target_override else args.target))
                send_heartbeat(update)

            ec = strax_proc.exitcode
            if ec is None:
                if t0 < now(-timeouts['max_processing_time']):
                    fail(f"Processing took longer than {timeouts['max_processing_time']} sec")
                    kill_process(strax_proc.pid)
                # Still working, check in later
                # TODO: is there a good way to detect hangs, before max_processing_time expires?

                log.info(f"Still processing run {run_id}")
                if args.production:
                    set_run_state(rd, 'busy', **info)
                time.sleep(timeouts['check_on_strax'])
                continue

            elif ec == 0:
                log.info(f"Strax done on run {run_id}, performing basic data quality check")

                try:
                    # Sometimes we have only he channels or mv channels, try loading one
                    # until we get one with chunks.
                    for rr_type in ('raw_records', 'raw_records_he', 'raw_records_mv',
                            # 'raw_records_nv' # not in the DAQ-reader yet.
                                    ):
                        md = st.get_meta(run_id, rr_type)
                        if len(md['chunks']) and (
                                'first_time' in md['chunks'][0] and
                                'last_endtime' in md['chunks'][0]
                        ):
                            break
                except Exception:
                    fail("Processing succeeded, but metadata not readable", error_traceback=traceback.format_exc())
                if not len(md['chunks']):
                    fail("Processing succeeded, but no chunks were written!")

                rd = get_run(mongo_id=rd['_id'])
                if 'end' not in rd:
                    fail("Processing succeeded, but run hasn't yet ended!")

                # Check that the data written covers the run
                # (at least up to some fudge factor)
                # Since chunks can be empty, and we don't want to crash,
                # this has to be done with some care...
                # Lets assume some ridiculous timestamp (in ns): 10e9*1e9
                t_covered = timedelta(
                    seconds=(max([x.get('last_endtime', 0) for x in md['chunks']]) -
                             min([x.get('first_time', 10e9*1e9) for x in md['chunks']])) / 1e9)
                run_duration = rd['end'] - rd['start']
                if not (0 < t_covered.seconds < float('inf')):
                    fail(f"Processed data covers {t_covered} sec")
                if not (timedelta(seconds=-max_timetamp_diff)
                        < (run_duration - t_covered)
                        < timedelta(seconds=max_timetamp_diff)):
                    fail(f"Processing covered {t_covered.seconds}, "
                         f"but run lasted {run_duration.seconds}!")
                if not all_files_saved(rd):
                    fail("Not all files in the rundoc for this run are saved")

                log.info(f"Run {run_id} processed successfully")
                if args.production:
                    set_run_state(rd, 'done', **info)
                    set_status_finished(rd)
                    upload_file_metadata(rd)

                    if args.delete_live:
                        delete_live_data(rd, loc)
                break

            else:
                # This is just the info that we're starting
                # exception retrieval. The actual error comes later.
                log.info(f"Failure while processing run {run_id}")
                if osp.exists(exception_tempfile):
                    with open(exception_tempfile, mode='r') as f:
                        exc_info = f.read()
                    if not exc_info:
                        exc_info = '[No exception info known, exception file was empty?!]'
                else:
                    exc_info = "[No exception info known, exception file not found?!]"
                fail(f"Strax exited with exit code {ec}.", error_traceback=f'Exception info: {exc_info}')
    except RunFailed:
        return


##
# Cleanup
##


def clean_run(*, mongo_id=None, number=None, force=False):
    """Removes all data on this host associated with a run
    that was previously registered in the run db.

    Does NOT remove temporary folders,
    nor data that isn't registered to the run db.
    """
    # Check mongo connection
    ping_dbs()

    # We need to get the full data docs here, since I was too lazy to write
    # a surgical update below
    rd = get_run(mongo_id=mongo_id, number=number, full_doc=True)
    have_live_data = False
    for dd in rd['data']:
        if dd['type'] == 'live':
            have_live_data = True
            break
    for ddoc in rd['data']:
        if 'host' in ddoc and ddoc['host'] == hostname:
            loc = ddoc['location']
            if not force and not have_live_data and 'raw_records' in ddoc['type']:
                log.info(f'prevent {loc} from being deleted. The live_data has already'
                         f' been removed')
            elif os.path.exists(loc):
                log.info(f'delete data at {loc}')
                _delete_data(rd, loc, ddoc['type'])
            else:
                loc = loc + '_temp'
                log.info(f'delete data at {loc}')
                _delete_data(rd, loc, ddoc['type'])


def clean_run_test_data(run_id):
    """
    Clean the data in the test_data_folder associated with this run_id
    """
    for folder in os.listdir(test_data_folder):
        if run_id in folder:
            log.info(f'Cleaning {test_data_folder + folder}')
            shutil.rmtree(test_data_folder + folder)


def cleanup_db():
    """Find various pathological runs and clean them from the db

    Also cleans the bootstrax collection for stale entries
    """
    # Check mongo connection
    ping_dbs()

    log.info("Checking for bad stuff in database")

    # Check for all the ebs if their last state message is not longer ago than the time we assume that the eb is dead.
    for eb_i in range(6):
        bd = bs_coll.find_one(
                {'host': f'eb{eb_i}.xenon.local'},
                sort=[('time', pymongo.DESCENDING)])
        if (bd and
                bd['time'].replace(tzinfo=pytz.utc) < now(
                    -timeouts['bootstrax_presumed_dead']) and
                bd['state'] is not dead_state):
            bs_coll.find_one_and_update({'_id': bd['_id']},
                                        {'$set': {'state': dead_state}})

    # Runs that say they are 'considering' or 'busy' but nothing happened for a while
    for state, timeout in [
        ('considering', timeouts['max_considering_time']),
        ('busy', timeouts['max_busy_time'])]:
        while True:
            send_heartbeat()
            rd = consider_run(
                {'bootstrax.state': state,
                 'bootstrax.time': {'$lt': now(-timeout)}},
                return_new_doc=False)
            if rd is None:
                break
            fail_run(rd,
                     f"Host {rd['bootstrax']['host']} said it was {state} "
                     f"at {rd['bootstrax']['time']}, but then didn't get further; "
                     f"perhaps it crashed on this run or is still stuck?")

    # Runs for which, based on the run doc alone, we can tell they are in a bad state
    # Mark them as failed.
    failure_queries = [
        ({'bootstrax.state': 'done',
          'end': {
              '$exists': False}},
         'Bootstrax state was done, but run did not yet end'),

        ({'bootstrax.state': 'done',
          'data': {
              '$not': {
                  '$elemMatch': {
                      "type": {
                          '$ne': 'live'}}}}},
         'Bootstrax state was done, but no processed data registered'),

        #       Can't add this yet, since registering happens when processing starts at the moment...
        #         ({'$not': {
        #             'bootstrax.state': 'done'},
        #           'data': {
        #               '$not': {
        #                   '$elemMatch': {
        #                       "type": {
        #                           '$ne': 'live'}}}}},
        #          'Bootstrax state was NOT done, but live data has been registered'),

        # For some reason this one doesn't work... probably I have my mongo query syntax confused
        #         ({'$and': [
        #             {'bootstrax.state': {
        #                 '$exists': True}},
        #             {'bootstrax.state': {
        #                 '$nin': 'considering done failed abandoned'.split()}}]},
        #          'Bootstrax state set to unrecognized state {bootstrax[state]}')
    ]

    for query, failure_message in failure_queries:
        while True:
            send_heartbeat()
            rd = consider_run(query)
            if rd is None:
                break
            fail_run(rd, failure_message.format(**rd))

    # Abandon runs which we already know are so bad that
    # there is no point in retrying them
    abandon_queries = [
        ({'tags': {
            '$elemMatch': {
                'name': 'bad'}}},
         "Run has a 'bad' tag"),
        ]

    for query, failure_message in abandon_queries:
        query['bootstrax.state'] = {'$ne': 'abandoned'}
        failure_message += ' -- run has been abandoned'
        while True:
            send_heartbeat()
            rd = consider_run(query)
            if rd is None:
                break
            fail_run(rd, failure_message.format(**rd))
            abandon(mongo_id=rd['_id'])


if __name__ == '__main__':
    if not args.undying:
        main()
    else:
        while True:
            try:
                main()
            except (KeyboardInterrupt, SystemExit):
                raise
            except Exception as fatal_error:
                log.error(f'Fatal warning:\tran into {fatal_error}. Try '
                          f'logging error and restart bootstrax')
                try:
                    log_warning(f'Fatal warning:\tran into {fatal_error}',
                                priority='error')
                except Exception as warning_error:
                    log.error(f'Fatal warning:\tcould not log {warning_error}')
                # This usually only takes a minute or two
                time.sleep(60)
                log.warning('Restarting main loop')
