bcache-loader

#!/usr/bin/env python3

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

import yaml

__version__ = "0.5.0"

# based on https://stackoverflow.com/a/42865957/2002471 adapted to our case
units = {"B": 1, "KB": 2**10, "MB": 2**20, "GB": 2**30, "TB": 2**40}


def parse_human_readable_size(size):
    size = size.upper()
    if not re.match(r' ', size):
        size = re.sub(r'([KMGTB]+)', r' \1', size)
    number, unit = [string.strip() for string in size.split()]
    return int(float(number)*units[unit + "B" if not unit.endswith('B') else ''])


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, retry=False):
    cache_basename = os.path.basename(c)
    cache_set_path = f'/sys/class/block/{cache_basename}/bcache/set'
    retries = 1 if retry else 4
    cache_set_found = False
    while (retries < 5):
        if os.path.exists(cache_set_path):
            cache_set_found = True
            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)) if cache_set_found else None


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 is not cache there's a chance that we are trying to re-register
        # the device because it changed size
        if not device_already_registered(d) or (dtype != "cache" and os.getenv('DEVICE_CHANGED')):
            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)
            retries = 1
            while (cmd_run.returncode):
                sl.syslog(sl.LOG_WARNING, f'Error while executing \'{" ".join(cmd)}\': {str(cmd_run.stderr)}')
                if (retries > 3):
                    break
                sleep(3)
                cmd_run = subprocess.run(cmd, capture_output=True)
                retries += 1
            return cmd_run.returncode
        else:
            sl.syslog(sl.LOG_INFO, f'Device {d} already registered')
    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)
    cache_set_path = f'/sys/class/block/{backing_basename}/bcache/cache'
    if os.path.exists(cache_set_path):
        cache_set = os.path.basename(os.readlink(cache_set_path))
        if cache_set == cset:
            sl.syslog(sl.LOG_INFO, f'{bd} already attached to {cset}.')
            return 0
        else:
            sl.syslog(sl.LOG_WARNING, f'{bd} is attached to a different cache: {cache_set} != {cset}. Trying to change it.')

    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 set_cutoff(bd, seq_cut):
    backing_basename = os.path.basename(bd)
    try:
        with open(f'/sys/class/block/{backing_basename}/bcache/sequential_cutoff', 'w') as f:
            f.write(str(seq_cut))
        return 0
    except Exception as e:
        sl.syslog(sl.LOG_ERR, f'Unable to set in {bd} sequential_cutoff to {seq_cut}. 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}')

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

        cache_set_uuid = get_cache_uuid(cd, True)

    for b in bds:
        backing_device = b['device']
        if register_device(backing_device, "backing", b.get('cache_mode')):
            if not device_already_registered(backing_device) or os.getenv('DEVICE_CHANGED'):
                # Indeed something went wrong with this device
                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)

        # Set sequential_cutoff
        if b.get('sequential_cutoff', -1) >= 0:
            set_cutoff(backing_device, b['sequential_cutoff'])
    return 0


def create_swap_flash_disk(swap_flash_path):
    real_swap_flash_path = os.readlink(swap_flash_path)
    bcache_match = re.search(r'(bcache\d+)', real_swap_flash_path)
    if not bcache_match:
        sl.syslog(sl.LOG_ERR, f"Couldn't find bcache device from the flash disk: '{real_swap_flash_path}'")
        return 1

    bcache_fd_basename = bcache_match[1]

    # Check if the disk is already a valid swap
    cmd = ['/usr/sbin/blkid', f'/dev/{bcache_fd_basename}']
    cmd_run = subprocess.run(cmd, capture_output=True, text=True)
    if cmd_run.returncode:
        cmd = ['/usr/sbin/mkswap', f'/dev/{bcache_fd_basename}']
        cmd_run = subprocess.run(cmd, capture_output=True, text=True)
        if cmd_run.returncode:
            sl.syslog(sl.LOG_ERR, f'Error while executing \'{" ".join(cmd)}\': {str(cmd_run.stderr)}')
            return cmd_run.returncode

    uuid_m = re.search(r'UUID="?([a-zA-z0-9\-]*)"?\s?', cmd_run.stdout)
    if not uuid_m:
        sl.syslog(sl.LOG_ERR, f'Couldn\'t extract the UUID from the `mkswap`. stdout: `{str(cmd_run.stdout)}`  stderr: `{str(cmd_run.stderr)}`')
        return 1

    uuid = uuid_m[1]

    # We need to create this symbolic link manually
    if not os.path.exists(f'/dev/disk/by-uuid/{uuid}'):
        fd = os.open('/dev/disk/by-uuid/', os.O_RDONLY | os.O_DIRECTORY)
        os.symlink(f'../../{bcache_fd_basename}', f'/dev/disk/by-uuid/{uuid}', dir_fd=fd)
        os.close(fd)

    cmd = ['/usr/sbin/swapon', '-U', uuid]
    cmd_run = subprocess.run(cmd, capture_output=True, text=True)
    if cmd_run.returncode:
        sl.syslog(sl.LOG_ERR, f'Error while executing \'{" ".join(cmd)}\': {str(cmd_run.stderr)}')
        return cmd_run.returncode

    sl.syslog(sl.LOG_INFO, f'Created swap disk (UUID={uuid}) on top of flash disk /dev/{bcache_fd_basename}')
    return 0


def get_flash_disks(cache_basename):
    flash_disks = []
    for v in glob(f'/sys/class/block/{cache_basename}/bcache/set/volume*'):
        fd = {"path": v}
        v_size_path = f'{v}/size'
        with open(v_size_path) as f:
            fd["size"] = f.readline().strip()

        v_label_path = f'{v}/label'
        with open(v_label_path) as f:
            fd["label"] = f.readline().strip()

        flash_disks.append(fd)
    return flash_disks


def create_flash_disks(cd, flash_disks):
    cache_basename = os.path.basename(cd)
    swap_disk_created = False
    for flash_disk in flash_disks:
        try:
            flash_disk_already_exists_path = ""
            flash_disk_size = flash_disk.get('size')
            flash_disk_label = flash_disk.get('label', "")

            if not flash_disk_size:
                sl.syslog(sl.LOG_ERR, 'No `size` defined for flash disk. Check your config.')
                continue

            if flash_disk_label == "swap" and swap_disk_created:
                sl.syslog(sl.LOG_ERR, 'Just one swap flash disk per cache device is allowed. Skipping creation.')
                continue

            # Check if already exists
            current_flash_disks = get_flash_disks(cache_basename)
            for fd in current_flash_disks:
                if (parse_human_readable_size(fd["size"]) == parse_human_readable_size(flash_disk_size)
                   and fd["label"] == flash_disk_label):
                    flash_disk_already_exists_path = fd["path"]
                    break

            if not flash_disk_already_exists_path:
                with open(f'/sys/class/block/{cache_basename}/bcache/set/flash_vol_create', 'w') as f:
                    f.write(flash_disk_size)

                sl.syslog(sl.LOG_INFO, f'Created flash disk of {flash_disk_size} under {cd}')
                current_flash_disks = get_flash_disks(cache_basename)
                created_flash_disk = next(iter(sorted((fd["path"] for fd in current_flash_disks
                                                       if parse_human_readable_size(fd["size"]) == parse_human_readable_size(flash_disk_size)),
                                               key=lambda v_path: int(os.path.basename(v_path).lstrip("volume")), reverse=True)),
                                          None)
                if not created_flash_disk:
                    sl.syslog(sl.LOG_ERR, "Couldn't find the just created flash disk")
                    continue

                with open(f'{created_flash_disk}/label', 'w') as f:
                    f.write(flash_disk_label)
            else:
                created_flash_disk = flash_disk_already_exists_path
                sl.syslog(sl.LOG_WARNING, f'A flash disk of the same size ({flash_disk_size}) and label ({flash_disk_label}) '
                          f'already exists: {flash_disk_already_exists_path}. Skipping creation...')

            if flash_disk_label == "swap":
                if create_swap_flash_disk(created_flash_disk):
                    sl.syslog(sl.LOG_ERR, 'Error while creating swap flash disk.')
                else:
                    swap_disk_created = True

        except Exception as e:
            sl.syslog(sl.LOG_ERR, f'Error while creating flash disk for {cd}. Reason: {e}')
            return 1
    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'])
        sys_arg_real_path = os.path.realpath(sys.argv[1])

        # Check if it's a cache device
        if sys_arg_real_path == cache_device:
            sl.syslog(sl.LOG_INFO, f'Managing cache device: {str(sys.argv[1])} (real path: {sys_arg_real_path})')
            attach_backing_and_cache([{**b, 'device': os.path.realpath(b['device'])} for b in cache['backing_devices']], cache_device)
            if 'flash_disks' in cache:
                create_flash_disks(cache_device, cache['flash_disks'])
        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_arg_real_path == backing_device:
                    sl.syslog(sl.LOG_INFO, f'Managing backing device: {backing["device"]} (real path: {backing_device})')
                    attach_backing_and_cache([{**backing, 'device': os.path.realpath(backing['device'])}], cache_device)
except Exception as e:
    sl.syslog(sl.LOG_ERR, f'Reason: {str(e)}')
    exit(1)