bcache-loader

#!/usr/bin/env python3

from logging import exception
import subprocess
import sys
import os
import syslog as sl
from time import sleep

import yaml


def device_exists(d):
    return os.path.exists(d)


def device_already_registered(d):
    device_basename = os.path.basename(d)
    return os.path.exists(f"/sys/class/block/{device_basename}/bcache/")


def get_cache_uuid(c):
    cache_basename = os.path.basename(c)
    cache_set_path = f'/sys/class/block/{cache_basename}/bcache/set'
    retries = 1
    while (retries < 5):
        if os.path.exists(cache_set_path):
            break
        sl.syslog(sl.LOG_WARNING, f'Path "{cache_set_path}" does not exist. Retrying in 3s...')
        sleep(3)
        retries += 1
    return os.path.basename(os.readlink(cache_set_path))


def probe_cache_device(d):
    cmd = ['/lib/udev/bcache-register', d]
    cmd_run = subprocess.run(cmd, capture_output=True)
    if cmd_run.returncode:
        sl.syslog(sl.LOG_WARNING, f'Error while executing \'{" ".join(cmd)}\': {str(cmd_run.stderr)}')

    return cmd_run.returncode


def register_device(d, dtype, cmode):
    if device_exists(d):
        if dtype != 'cache' or not device_already_registered(d):
            cmd = ['/usr/sbin/make-bcache']
            if dtype == 'cache':
                cmd += ['-C', d]
            else:
                cmd += ['--ioctl', '-B', d] + (['--writeback'] if cmode and str.lower(cmode) == 'wb' else [])
            cmd_run = subprocess.run(cmd, capture_output=True)
            if cmd_run.returncode:
                sl.syslog(sl.LOG_WARNING, f'Error while executing \'{" ".join(cmd)}\': {str(cmd_run.stderr)}')

            return cmd_run.returncode
    else:
        sl.syslog(sl.LOG_WARNING, f'Device does not exist: {d}')

    return 0


def attach_backing_to_cache(bd, cset):
    backing_basename = os.path.basename(bd)
    try:
        with open(f'/sys/class/block/{backing_basename}/bcache/attach', 'w') as f:
            f.write(cset)
        return 0
    except Exception as e:
        sl.syslog(sl.LOG_ERR, f'Unable to attach {bd} to {cset}. Reason: {str(e)}')
        return 1


def attach_backing_and_cache(bds, cd):

    cache_set_uuid = None

    if register_device(cd, "cache", None):
        sl.syslog(sl.LOG_ERR, f'Error while registering cache device {cd}')
    else:
        sl.syslog(sl.LOG_INFO, f'Successfully registered cache device {cd}')

    if probe_cache_device(cd) == 0:
        sl.syslog(sl.LOG_INFO, f'Successfully probed cache device {cd}')

    cache_set_uuid = get_cache_uuid(cd)

    for b in bds:
        backing_device = b['device']
        if register_device(backing_device, "backing", b['cache_mode']):
            sl.syslog(sl.LOG_ERR, f'Error while registering backing device {backing_device} ...')
            return 1

        sl.syslog(sl.LOG_INFO, f'Successfully registered backing device {backing_device}')
        sleep(1)  # Wait for the backing device to fully register
        if cache_set_uuid:
            sl.syslog(sl.LOG_INFO, f'Attaching backing device {backing_device} to cache device {cd} with UUID {cache_set_uuid}')
            attach_backing_to_cache(backing_device, cache_set_uuid)

    return 0


try:
    subprocess.call(['/sbin/modprobe', 'bcache'])
except Exception as e:
    sl.syslog(sl.LOG_ERR, f'Unable to probe custom_bcache module. Reason: {str(e)}')
    exit(1)

try:
    with open('/etc/bcache/bcache.conf') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
except Exception as e:
    sl.syslog(sl.LOG_ERR, f'Unable to load bcache config. Reason: {str(e)}')
    exit(1)

try:

    for cache in config['cache_devices']:
        cache_device = os.path.realpath(cache['device'])

        # Check if it's a cache device
        if sys.argv[1] == cache_device:
            sl.syslog(sl.LOG_INFO, f'Managing cache device: {str(sys.argv[1])}')
            attach_backing_and_cache([os.path.realpath(b) for b in cache['backing_devices']], cache_device)
        else:
            # Check if it's a backing device of this cache device
            for backing in cache['backing_devices']:
                backing_device = os.path.realpath(backing['device'])
                if sys.argv[1] == backing_device:
                    sl.syslog(sl.LOG_INFO, f'Managing backing device: {str(sys.argv[1])}')
                    attach_backing_and_cache([backing], cache_device)
except Exception as e:
    sl.syslog(sl.LOG_ERR, f'Reason: {str(e)}')
    exit(1)