Python module for using Linux's futex() syscall, and in particular the FUTEX_WAIT_MULTIPLE operation.

fwaitm.py

#!/usr/bin/env python3
# Copyright (c) 2020 Torge Matthies
#
# This software is provided 'as-is', without any express or implied
# warranty. In no event will the authors be held liable for any damages
# arising from the use of this software.
#
# Permission is granted to anyone to use this software for any purpose,
# including commercial applications, and to alter it and redistribute it
# freely, subject to the following restrictions:
#
# 1. The origin of this software must not be misrepresented; you must not
#    claim that you wrote the original software. If you use this software
#    in a product, an acknowledgment in the product documentation would be
#    appreciated but is not required.
# 2. Altered source versions must be plainly marked as such, and must not be
#    misrepresented as being the original software.
# 3. This notice may not be removed or altered from any source distribution.
#
# Author contact info:
#    E-Mail address: openglfreak@googlemail.com
#    PGP key fingerprint: 0535 3830 2F11 C888 9032 FAD2 7C95 CD70 C9E8 438D

'''Module for using Linux's futex() syscall, and in particular the
FUTEX_WAIT_MULTIPLE operation.
'''
# pylint: disable=invalid-name,too-few-public-methods
import ctypes
import errno
import os
import subprocess

###############
# futex stuff #
###############

__all__ = ['timespec', 'get_futex_syscall_nr', 'get_futex_syscall']


class timespec(ctypes.Structure):
    '''Linux kernel compatible timespec type.

    Fields:
        tv_sec: The whole seconds of the timespec.
        tv_nsec: The nanoseconds of the timespec.
    '''
    __slots__ = ()
    _fields_ = [
        ('tv_sec', ctypes.c_long),
        ('tv_nsec', ctypes.c_long),
    ]


# Hardcode some of the most commonly used architectures's
# futex syscall numbers.
_NR_FUTEX_PER_ARCH = {
    ('i386', 32): 240,
    ('i686', 32): 240,
    ('x86_64', 32): 240,
    ('x86_64', 64): 202,
    ('aarch64', 64): 240,
    ('aarch64_be', 64): 240,
    ('armv8b', 32): 240,
    ('armv8l', 32): 240,
}


def _get_futex_syscall_nr():
    bits = ctypes.sizeof(ctypes.c_void_p) * 8

    try:
        return _NR_FUTEX_PER_ARCH[(os.uname()[4], bits)]
    except KeyError:
        pass

    try:
        with subprocess.Popen(
                ('cpp', '-m' + str(bits), '-E', '-P', '-x', 'c', '-'),
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                close_fds=True,
                universal_newlines=True,
        ) as popen:
            stdout, stderr = popen.communicate(
                '#include <sys/syscall.h>\n'
                '__NR_futex\n'
            )
    except FileNotFoundError as ex:
        raise RuntimeError(
            'failed to determine futex syscall number: '
            'cpp not installed or not in PATH'
        ) from ex

    if popen.returncode:
        raise RuntimeError(
            'failed to determine futex syscall number: '
            'cpp returned nonzero exit code',
            stderr
        )

    if not stdout:
        raise RuntimeError(
            'failed to determine futex syscall number: '
            'no output from cpp'
        )

    last_line = stdout.splitlines()[-1]

    if last_line == '__NR_futex':
        raise RuntimeError(
            'failed to determine futex syscall number: '
            '__NR_futex not expanded'
        )

    try:
        return int(last_line)
    except ValueError as ex:
        raise RuntimeError(
            'failed to determine futex syscall number: '
            '__NR_futex not a valid number: ' + last_line
        ) from ex

    assert False


_CACHED_FUTEX_SYSCALL_NR = None


def get_futex_syscall_nr():
    '''Get the syscall number of the Linux futex() syscall.

    Returns:
        The futex() syscall number.

    Raises:
        RuntimeError: When the syscall number could not be determined.
    '''
    global _CACHED_FUTEX_SYSCALL_NR  # pylint: disable=global-statement
    if _CACHED_FUTEX_SYSCALL_NR is None:
        _CACHED_FUTEX_SYSCALL_NR = _get_futex_syscall_nr()

        def get_futex_syscall_nr_cached():
            return _CACHED_FUTEX_SYSCALL_NR
        global get_futex_syscall_nr  # pylint: disable=global-variable-undefined # noqa: E501
        get_futex_syscall_nr_cached.__doc__ = \
            getattr(get_futex_syscall_nr, '__doc__', None)
        get_futex_syscall_nr = get_futex_syscall_nr_cached

    return _CACHED_FUTEX_SYSCALL_NR


def _is_ctypes_obj(obj):
    return (
        hasattr(obj, '_b_base_')
        and hasattr(obj, '_b_needsfree_')
        and hasattr(obj, '_objects')
    )


def _is_ctypes_obj_pointer(obj):
    return hasattr(obj, '_type_') and hasattr(obj, 'contents')


def _coerce_to_pointer(obj):
    if obj is None:
        return None

    if _is_ctypes_obj(obj):
        if _is_ctypes_obj_pointer(obj):
            return obj
        return ctypes.pointer(obj)

    obj = tuple(obj)
    return (obj[0].__class__ * len(obj))(*obj)


def _get_futex_syscall():
    futex_syscall = ctypes.CDLL(None, use_errno=True).syscall
    futex_syscall.argtypes = (ctypes.c_long, ctypes.c_void_p, ctypes.c_int,
                              ctypes.c_int, ctypes.POINTER(timespec),
                              ctypes.c_void_p, ctypes.c_int)
    futex_syscall.restype = ctypes.c_int
    futex_syscall_nr = get_futex_syscall_nr()

    # pylint: disable=too-many-arguments
    def _futex_syscall(uaddr, futex_op, val, timeout, uaddr2, val3):
        error = futex_syscall(
            futex_syscall_nr,
            _coerce_to_pointer(uaddr),
            futex_op,
            val,
            _coerce_to_pointer(timeout or timespec()),
            _coerce_to_pointer(uaddr2),
            val3
        )
        return error, (ctypes.get_errno() if error == -1 else 0)
    _futex_syscall.__doc__ = getattr(futex, '__doc__', None)

    return _futex_syscall


_CACHED_FUTEX_SYSCALL = None


def get_futex_syscall():
    '''Create a function that can be used to execute the Linux futex()
    syscall.

    Returns:
        A proxy function for the Linux futex() syscall.

    Raises:
        AttributeError: When the libc has no syscall() function.
        RuntimeError: When the syscall number could not be determined.
    '''
    global _CACHED_FUTEX_SYSCALL  # pylint: disable=global-statement
    if _CACHED_FUTEX_SYSCALL is None:
        _CACHED_FUTEX_SYSCALL = _get_futex_syscall()

        def get_futex_syscall_cached():
            return _CACHED_FUTEX_SYSCALL
        global get_futex_syscall  # pylint: disable=global-variable-undefined
        get_futex_syscall_cached.__doc__ = \
            getattr(get_futex_syscall, '__doc__', None)
        get_futex_syscall = get_futex_syscall_cached

    return _CACHED_FUTEX_SYSCALL


# pylint: disable=too-many-arguments
def futex(uaddr, futex_op, val, timeout, uaddr2, val3):
    '''Invoke the Linux futex() syscall with the provided arguments.

    Args:
        See the description of the futex() syscall for the parameter
        meanings.
        `uaddr` and `uaddr2` are automatically converted to pointers.
        If timeout is None, a zero timeout is passed.

    Returns:
        A tuple of the return value of the syscall and the error code
        in case an error occurred.

    Raises:
        AttributeError: When the libc has no syscall() function.
        RuntimeError: When the syscall number could not be determined.
        TypeError: If `uaddr` or `uaddr2` is not a pointer and can't be
            converted into one.
    '''
    global _CACHED_FUTEX_SYSCALL  # pylint: disable=global-statement
    futex_syscall = _CACHED_FUTEX_SYSCALL
    if futex_syscall is None:
        futex_syscall = get_futex_syscall()

        global futex  # pylint: disable=global-variable-undefined
        futex = futex_syscall

    return futex_syscall(uaddr, futex_op, val, timeout, uaddr2, val3)


#################################
# FUTEX_WAIT_MULTIPLE functions #
#################################

__all__.extend(['get_futex_wait_multiple_op',
                'get_futex_wait_multiple_syscall', 'futex_wait_multiple',
                'is_futex_wait_multiple_supported'])


def _get_futex_wait_multiple_op():
    futex_syscall = get_futex_syscall()
    ret = futex_syscall(None, 31, 0, None, None, 0)
    if ret[1] != errno.ENOSYS:
        return 31
    ret = futex_syscall(None, 13, 0, None, None, 0)
    if ret[1] != errno.ENOSYS:
        return 13
    return None


_CACHED_FUTEX_WAIT_MULTIPLE_OP = None


def get_futex_wait_multiple_op():
    '''Get the operation code of the Linux futex FUTEX_WAIT_MULTIPLE
    operation.

    Returns:
        The operation code of the FUTEX_WAIT_MULTIPLE operation.

    Raises:
        AttributeError: When the libc has no syscall() function.
        RuntimeError: When the Linux futex() syscall number could not be
            determined.
    '''
    global _CACHED_FUTEX_WAIT_MULTIPLE_OP  # pylint: disable=global-statement
    if _CACHED_FUTEX_WAIT_MULTIPLE_OP is None:
        _CACHED_FUTEX_WAIT_MULTIPLE_OP = _get_futex_wait_multiple_op()

        def get_futex_wait_multiple_op_cached():
            return _CACHED_FUTEX_WAIT_MULTIPLE_OP
        global get_futex_wait_multiple_op  # pylint: disable=global-variable-undefined # noqa: E501
        get_futex_wait_multiple_op_cached.__doc__ = \
            getattr(get_futex_wait_multiple_op, '__doc__', None)
        get_futex_wait_multiple_op = get_futex_wait_multiple_op_cached

    return _CACHED_FUTEX_WAIT_MULTIPLE_OP


def _get_futex_wait_multiple_syscall():
    futex_syscall = get_futex_syscall()
    futex_wait_multiple_op = get_futex_wait_multiple_op()
    if futex_wait_multiple_op is None:
        return lambda futexes, timeout: (-1, errno.ENOSYS)

    def _futex_wait_multiple(futexes, timeout):
        return futex_syscall(
            futexes,
            futex_wait_multiple_op,
            len(futexes) if futexes is not None else 0,
            timeout,
            None,
            0
        )
    _futex_wait_multiple.__doc__ = \
        getattr(futex_wait_multiple, '__doc__', None)

    return _futex_wait_multiple


_CACHED_FUTEX_WAIT_MULTIPLE_SYSCALL = None


def get_futex_wait_multiple_syscall():
    '''Create a function that can be used to invoke the Linux futex
    FUTEX_WAIT_MULTIPLE operation.

    Returns:
        A proxy function for the Linux futex FUTEX_WAIT_MULTIPLE
        operation.

    Raises:
        AttributeError: When the libc has no syscall() function.
        RuntimeError: When the Linux futex() syscall number could not be
            determined.
    '''
    global _CACHED_FUTEX_WAIT_MULTIPLE_SYSCALL  # pylint: disable=global-statement # noqa: E501
    if _CACHED_FUTEX_WAIT_MULTIPLE_SYSCALL is None:
        _CACHED_FUTEX_WAIT_MULTIPLE_SYSCALL = \
            _get_futex_wait_multiple_syscall()

        def get_futex_wait_multiple_syscall_cached():
            return _CACHED_FUTEX_WAIT_MULTIPLE_SYSCALL
        global get_futex_wait_multiple_syscall  # pylint: disable=global-variable-undefined # noqa: E501
        get_futex_wait_multiple_syscall_cached.__doc__ = \
            getattr(get_futex_wait_multiple_syscall, '__doc__', None)
        get_futex_wait_multiple_syscall = \
            get_futex_wait_multiple_syscall_cached

    return _CACHED_FUTEX_WAIT_MULTIPLE_SYSCALL


def futex_wait_multiple(futexes, timeout):
    '''Invoke the Linux futex FUTEX_WAIT_MULTIPLE operation with the
    provided arguments.

    Args:
        futexes: An iterable of futex_wait_block instances.
        timeout: The timeout for the operation. May be None for a zero
            timeout.

    Returns:
        A tuple of the return value of the operation and the error code
        in case an error occurred.

    Raises:
        AttributeError: When the libc has no syscall() function.
        RuntimeError: When the syscall number could not be determined.
        TypeError: If `futexes` or `timeout` has a wrong type.
    '''
    global _CACHED_FUTEX_WAIT_MULTIPLE_SYSCALL  # pylint: disable=global-statement # noqa: E501
    futex_wait_multiple_syscall = _CACHED_FUTEX_WAIT_MULTIPLE_SYSCALL
    if futex_wait_multiple_syscall is None:
        futex_wait_multiple_syscall = get_futex_wait_multiple_syscall()

        global futex_wait_multiple  # pylint: disable=global-variable-undefined
        futex_wait_multiple = futex_wait_multiple_syscall

    return futex_wait_multiple_syscall(futexes, timeout)


def _is_futex_wait_multiple_supported():
    try:
        return futex_wait_multiple(None, None)[1] != errno.ENOSYS
    except (AttributeError, RuntimeError):
        return False


_CACHED_FUTEX_WAIT_MULTIPLE_SUPPORTED = None


def is_futex_wait_multiple_supported():
    '''Checks whether the Linux futex FUTEX_WAIT_MULTIPLE operation is
    supported on this kernel.

    Returns:
        Whether this kernel supports the FUTEX_WAIT_MULTIPLE operation.
    '''
    global _CACHED_FUTEX_WAIT_MULTIPLE_SUPPORTED  # pylint: disable=global-statement # noqa: E501
    futex_wait_multiple_supported = _CACHED_FUTEX_WAIT_MULTIPLE_SUPPORTED
    if futex_wait_multiple_supported is None:
        futex_wait_multiple_supported = _is_futex_wait_multiple_supported()

        def is_futex_wait_multiple_supported_cached():
            return futex_wait_multiple_supported
        global is_futex_wait_multiple_supported  # pylint: disable=global-variable-undefined # noqa: E501
        is_futex_wait_multiple_supported_cached.__doc__ = \
            getattr(is_futex_wait_multiple_supported, '__doc__', None)
        is_futex_wait_multiple_supported = \
            is_futex_wait_multiple_supported_cached

    return futex_wait_multiple_supported


#############################
# FUTEX_WAIT_MULTIPLE types #
#############################

__all__.extend(['futex_type', 'futex_pointer_type', 'futex_wait_block'])


class futex_type(ctypes.c_int32):
    '''Futex type (32-bit integer).'''
    __slots__ = ()


class futex_pointer_type(ctypes.POINTER(futex_type)):
    '''Pointer to a futex.'''
    __slots__ = ()


class futex_wait_block:
    '''Futex specification structure used in calls to the
    FUTEX_WAIT_MULTIPLE futex operation.

    Fields:
        addr: Pointer to a futex.
        val: The value to compare the futex against.
        bitset: Bitset to compare against in the FUTEX_WAKE_BITSET
            operation.
    '''
    __slots__ = ()

    def __new__(cls, *args, **kwargs):
        futex_syscall = get_futex_syscall()
        futex_wait_multiple_op = get_futex_wait_multiple_op()
        if futex_wait_multiple_op is None:
            raise TypeError('FUTEX_WAIT_MULTIPLE not supported')

        # Try _futex_wait_block_native.
        fwb = _futex_wait_block_native(
            futex_pointer_type(futex_type(0)),
            0,
            0xFFFFFFFF
        )
        err = futex_syscall(
            (fwb, _futex_wait_block_native()),
            futex_wait_multiple_op,
            1,
            None,
            None,
            0
        )[1]
        if err == errno.ETIMEDOUT:
            def __new__cached(_, *_args, **_kwargs):
                return _futex_wait_block_native(*_args, **_kwargs)
            cls.__new__ = __new__cached
            return _futex_wait_block_native(*args, **kwargs)

        # Try _futex_wait_block_32on64.
        fwb = _futex_wait_block_32on64(
            futex_pointer_type(futex_type(0)),
            0,
            0xFFFFFFFF
        )
        err = futex_syscall(
            (fwb, _futex_wait_block_32on64()),
            futex_wait_multiple_op,
            1,
            None,
            None,
            0
        )[1]
        if err == errno.ETIMEDOUT:
            def __new__cached(_, *_args, **_kwargs):
                return _futex_wait_block_32on64(*_args, **_kwargs)
            cls.__new__ = __new__cached
            return _futex_wait_block_32on64(*args, **kwargs)

        # Neither worked.
        raise TypeError(
            'Could not determine which type of futex_wait_block to use'
        )

    @staticmethod
    def __instancecheck__(instance):
        return isinstance(
            instance,
            (_futex_wait_block_native, _futex_wait_block_32on64)
        )


class _futex_wait_block_native(ctypes.Structure):
    __doc__ = getattr(futex_wait_block, '__doc__', None)

    __metaclass__ = futex_wait_block
    __slots__ = ()
    _fields_ = [
        ('addr', futex_pointer_type),
        ('val', futex_type),
        ('bitset', ctypes.c_int32),
    ]


class _futex_wait_block_32on64(ctypes.Structure):
    __doc__ = getattr(futex_wait_block, '__doc__', None)

    class _FutexPointer64(ctypes.Union):
        __slots__ = ()
        _fields_ = [
            ('addr', futex_pointer_type),
            ('_dummy', ctypes.c_uint64),
        ]

    __metaclass__ = futex_wait_block
    __slots__ = ()
    _anonymous_ = ('_addr',)
    _fields_ = [
        ('_addr', _FutexPointer64),
        ('val', futex_type),
        ('bitset', ctypes.c_int32),
    ]

    def __init__(self, *args, **kwargs):
        if not args:
            super().__init__(**kwargs)
            return
        if not isinstance(args[0], _futex_wait_block_32on64._FutexPointer64):
            args = _futex_wait_block_32on64._FutexPointer64(args[0]), *args[1:]
        super().__init__(*args, **kwargs)


###############
# Main stuff. #
###############

__all__.extend(['print_help', 'main'])

_HELP_TEXT = '''Usage: %s [options]

options:
    -h / --help:  show this help text
    -q / --quiet: do not output result

exit status is 0 if FUTEX_WAIT_MULTIPLE is supported, 1 if not,
and 2 for an unknown option.'''


def print_help(arg0):
    '''Print the command-line help text for this module.

    Args:
        arg0: The name of this program.
    '''
    print(_HELP_TEXT % (arg0))


def main(argv):
    print_result = True

    for arg in argv[1:]:
        if arg in ('-h', '--help'):
            print_help(argv[0])
            return 0
        if arg in ('-q', '--quiet'):
            print_result = False
        else:
            print('Unknown parameter:', arg, file=sys.stderr)
            return 2

    if is_futex_wait_multiple_supported():
        if print_result:
            print('FUTEX_WAIT_MULTIPLE is supported.')
        return 0
    if print_result:
        print('FUTEX_WAIT_MULTIPLE is not supported.')
    return 1


if __name__ == '__main__':
    import sys

    raise SystemExit(main(sys.argv))