#!/usr/bin/python3
# Copyright 2017 Red Hat, Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
#
# Refer to the README and COPYING files for full details of the license
#

import argparse
import ctypes
import errno
import io
import mmap
import os

from contextlib import closing


def main():
    options = parse_args()

    # Use direct I/O so writing zeroes do not delay other I/O on the system.
    fd = os.open(options.filename, os.O_RDWR | os.O_CREAT | os.O_DIRECT)
    try:
        # First try native fallocate. On NFS 4.2, GlusterFS, XFS and ext4, this
        # is practically free.
        native_fallocate(fd, options.offset, options.size)
    except OSError as e:
        if e.errno != errno.EOPNOTSUPP:
            raise

        # On NFS < 4.2 fallback to writing zeroes. Surprisingly this is 2.5
        # times faster than posix_fallocate() and use 12 times less cpu time.
        write_zeroes(fd, options.offset, options.size)
    finally:
        os.close(fd)


def parse_args():
    parser = argparse.ArgumentParser(
        description='fallocate is used to preallocate blocks to a file.')

    parser.add_argument(
        '--offset',
        type=size,
        default=0,
        help='Offset in bytes to start allocation from. Use t|g|m|k to '
             'specify size in terabytes, gigabytes, megabytes or '
             'kilobytes')

    parser.add_argument(
        'size',
        type=size,
        help='Size in bytes to allocate. Use t|g|m|k to specify size in '
             'terabytes, gigabytes, megabytes or kilobytes')

    parser.add_argument(
        'filename',
        help='Name of file to allocate')

    return parser.parse_args()


def size(s):
    if s.endswith("t"):
        return int(s[:-1]) * 1024**4
    elif s.endswith("g"):
        return int(s[:-1]) * 1024**3
    elif s.endswith("m"):
        return int(s[:-1]) * 1024**2
    elif s.endswith("k"):
        return int(s[:-1]) * 1024**1
    else:
        return int(s)


def native_fallocate(fd, offset, length):
    libc = ctypes.CDLL("libc.so.6", use_errno=True)

    if not hasattr(libc, "fallocate"):
        raise os_error(errno.EOPNOTSUPP)

    err = libc.fallocate(
        fd, 0, ctypes.c_longlong(offset), ctypes.c_longlong(length))
    if err != 0:
        raise os_error(ctypes.get_errno())


def write_zeroes(fd, offset, length, buffer_size=8 * 1024**2):
    """
    Allocate file space by writing zeroes.

    Note: offset and size must be aligned to underlying storage logical block
    size.
    """
    buffer_size = min(buffer_size, length)

    buf = mmap.mmap(-1, buffer_size, mmap.MAP_SHARED)
    with closing(buf):
        # fd is owned by caller.
        f = io.FileIO(fd, "r+", closefd=False)
        with closing(f):
            f.seek(offset)

            # Write complete buffers.
            while length > len(buf):
                try:
                    length -= f.write(buf)
                except InterruptedError:
                    pass

            # Write last buffer.
            while length > 0:
                with memoryview(buf)[:length] as v:
                    try:
                        length -= f.write(v)
                    except InterruptedError:
                        pass

            os.fsync(fd)


def os_error(err):
    return OSError(err, os.strerror(err))


if __name__ == '__main__':
    main()
