btrfs-inspect-parser.py

#!/usr/bin/env python3
"""
This is a quick and dirty script to find btrfs blocks that have certain properties.
It does this by parsing the output produced by "btrfs inspect-internal dump-tree"
and outputs virtual byte ranges as used by the vrange option of "btrfs balance start"

I needed this to re-balance blocks that were not on a specific device of my btrfs
raid1 after it became full and I added a new drive. If you need it for something
else your mileage may vary.

This code is released into the public domain. Use at your own risk.

edit the code in 'if __name__ == "__main__":' as needed

usage:
sudo btrfs inspect-internal dump-tree --hide-names /dev/mapper/some_btrfs_raid_member | btrfs-inspect-parser.py | sudo xargs -I{} btrfs balance start -dvrange={} /mnt/some_btrfs_raid

--hide-names is here to avoid problems with weird file names
"""


BTRFS_BLOCK_SIZE = 1073741824


class IndentBlock:
    """
    Very stupid generic parser for indented text files
    """
    def __init__(self, name, data, children):
        self.name = name
        self.data = data.rstrip()
        self.children = children

    @classmethod
    def from_lines(cls, lines, first=None):
        """
        Parse one indented block from lines.
        
        ::args:
          lines: iterator of lines (str)
          first: first line (str)
        
        ::returns: block, next_line
          block will be a IndentBlock instance or None if iterator is empty
          next_line will be None if EOF else str (first line of next block)
        """
        if first is None:
            try:
                first = next(lines)
            except StopIteration:
                return None, None
        first_stripped = first.lstrip()
        init_level = len(first) - len(first_stripped)
        name, data = first_stripped.split(" ", maxsplit=1)
        children = []
        line = None
        while True:
            if line is None:
                try:
                    line = next(lines)
                except StopIteration:
                    break
            line_stripped = line.lstrip()
            current_level = len(line) - len(line_stripped)
            if current_level <= init_level:
                break
            else:
                child, line = cls.from_lines(lines, line)
                children.append(child)
        return cls(name, data, children), line
    
    def print(self, indent=0):
        print(" " * indent + self.name + " " + self.data)
        for child in self.children:
            child.print(indent+2)
    
    def __str__(self):
        return self.name + " " + self.data
    
    def __repr__(self):
        return f"<{self.__class__.__name__}> of {self.name} with {len(self.children)} children"


class StripedItem:
    def __init__(self, item_block):
        if not self.is_striped(item_block):
            raise Exception("not a striped item block")
        self.block_group = self.parse_block_group(item_block.data)
        self.block_type = self.parse_block_type(item_block.children)
        self.devices = self.parse_device_ids(item_block.children)
        self.block = item_block

    @staticmethod
    def is_striped(block):
        return block.name == "item" and " key (FIRST_CHUNK_TREE CHUNK_ITEM " in block.data

    @staticmethod
    def parse_block_group(data):
        return int(data.split(" ")[4].rstrip(")"))

    @staticmethod    
    def parse_device_ids(children):
        ids = set()
        for child in children:
            if child.name == "num_stripes":
                for s_child in child.children:
                    if s_child.name == "stripe":
                        ids.add(int(s_child.data.split(" ")[2]))
        return ids
    
    @staticmethod    
    def parse_block_type(children):
        for child in children:
            if child.name == "length":
                btype = child.data.rsplit(" ", maxsplit=1)[-1]
                return btype.split("|")
        return None, None

    def __repr__(self):
        return f"<{self.__class__.__name__}> of {self.block_group} on devid "\
               f"{','.join(map(str, self.devices))} ({'|'.join(self.block_type)})"


def iter_blocks(lines_iter, name=None):
    next_line = None
    while True:
        block, next_line = IndentBlock.from_lines(lines_iter, next_line)
        if not block:
            return
        if not name or block.name == name:
            yield block


def find_item_blocks(block):
    if block.name != "chunk":
        return
    for child in block.children:
        if StripedItem.is_striped(child):
            yield StripedItem(child)


def combine_offsets(offsets):
    if not offsets:
        return {}
    last_key = offsets[0]
    last = last_key
    combined = {last_key: 1}
    for off in offsets[1:]:
        if off != last + BTRFS_BLOCK_SIZE:
            combined[off] = 1
            last_key = off
        else:
            combined[last_key] += 1
        last = off
    return combined


if __name__ == "__main__":
    #fn = "/tmp/btrfs-inspect-internal-dump-tree-output"
    #with open(fn, "r") as fh:
    raise Exception("there is no support for this code. you're on your own.")
    import sys
    fh = sys.stdin
    device_id = 1
    i = 0
    block_offsets = []
    for block in iter_blocks(fh, "chunk"):
        for item in find_item_blocks(block):
            # find all DATA blocks that are not on device with id device_id
            if item.block_type[0] == 'DATA' and device_id not in item.devices:
                block_offsets.append(item.block_group)
    block_offsets.sort()
    for offset, count in combine_offsets(block_offsets).items():
        print(f"{offset}..{offset + count * BTRFS_BLOCK_SIZE}")