#!/usr/bin/env python3
# group: rw
#
# Test graph changes while I/O is happening
#
# Copyright (C) 2022 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, see <http://www.gnu.org/licenses/>.
#

import os
from threading import Thread
import iotests
from iotests import imgfmt, qemu_img, qemu_img_create, qemu_io, \
        QMPTestCase, QemuStorageDaemon


top = os.path.join(iotests.test_dir, 'top.img')
mid = os.path.join(iotests.test_dir, 'mid.img')
nbd_sock = os.path.join(iotests.sock_dir, 'nbd.sock')


def do_qemu_img_bench(count: int = 2000000) -> None:
    """
    Do some I/O requests on `nbd_sock`.
    """
    qemu_img('bench', '-f', 'raw', '-c', str(count),
             f'nbd+unix:///node0?socket={nbd_sock}')


class TestGraphChangesWhileIO(QMPTestCase):
    def setUp(self) -> None:
        # Create an overlay that can be added at runtime on top of the
        # null-co block node that will receive I/O
        qemu_img_create('-f', imgfmt, '-F', 'raw', '-b', 'null-co://', top)

        # QSD instance with a null-co block node in an I/O thread,
        # exported over NBD (on `nbd_sock`, export name "node0")
        self.qsd = QemuStorageDaemon(
            '--object', 'iothread,id=iothread0',
            '--blockdev', 'null-co,node-name=node0,read-zeroes=true',
            '--nbd-server', f'addr.type=unix,addr.path={nbd_sock}',
            '--export', 'nbd,id=exp0,node-name=node0,iothread=iothread0,' +
                        'fixed-iothread=true,writable=true',
            qmp=True
        )

    def tearDown(self) -> None:
        self.qsd.stop()
        os.remove(top)

    def _wait_for_blockjob(self, status: str) -> None:
        done = False
        while not done:
            for event in self.qsd.get_qmp().get_events(wait=10.0):
                if event['event'] != 'JOB_STATUS_CHANGE':
                    continue
                if event['data']['status'] == status:
                    done = True

    def test_blockdev_add_while_io(self) -> None:
        # Run qemu-img bench in the background
        bench_thr = Thread(target=do_qemu_img_bench)
        bench_thr.start()

        # While qemu-img bench is running, repeatedly add and remove an
        # overlay to/from node0
        while bench_thr.is_alive():
            self.qsd.cmd('blockdev-add', {
                'driver': imgfmt,
                'node-name': 'overlay',
                'backing': 'node0',
                'file': {
                    'driver': 'file',
                    'filename': top
                }
            })

            self.qsd.cmd('blockdev-del', {
                'node-name': 'overlay'
            })

        bench_thr.join()

    def test_commit_while_io(self) -> None:
        # Run qemu-img bench in the background
        bench_thr = Thread(target=do_qemu_img_bench, args=(200000, ))
        bench_thr.start()

        qemu_io('-c', 'write 0 64k', top)
        qemu_io('-c', 'write 128k 64k', top)

        self.qsd.cmd('blockdev-add', {
            'driver': imgfmt,
            'node-name': 'overlay',
            'backing': None,
            'file': {
                'driver': 'file',
                'filename': top
            }
        })

        self.qsd.cmd('blockdev-snapshot', {
            'node': 'node0',
            'overlay': 'overlay',
        })

        # While qemu-img bench is running, repeatedly commit overlay to node0
        while bench_thr.is_alive():
            self.qsd.cmd('block-commit', {
                'job-id': 'job0',
                'device': 'overlay',
            })

            self.qsd.cmd('block-job-cancel', {
                'device': 'job0',
            })

            self._wait_for_blockjob('null')

        bench_thr.join()

    def test_remove_lower_snapshot_while_io(self) -> None:
        # Run qemu-img bench in the background
        bench_thr = Thread(target=do_qemu_img_bench, args=(100000, ))
        bench_thr.start()

        # While I/O is performed on 'node0' node, consequently add 2 snapshots
        # on top of it, then remove (commit) them starting from lower one.
        while bench_thr.is_alive():
            # Recreate snapshot images on every iteration
            qemu_img_create('-f', imgfmt, mid, '1G')
            qemu_img_create('-f', imgfmt, top, '1G')

            self.qsd.cmd('blockdev-add', {
                'driver': imgfmt,
                'node-name': 'mid',
                'file': {
                    'driver': 'file',
                    'filename': mid
                }
            })

            self.qsd.cmd('blockdev-snapshot', {
                'node': 'node0',
                'overlay': 'mid',
            })

            self.qsd.cmd('blockdev-add', {
                'driver': imgfmt,
                'node-name': 'top',
                'file': {
                    'driver': 'file',
                    'filename': top
                }
            })

            self.qsd.cmd('blockdev-snapshot', {
                'node': 'mid',
                'overlay': 'top',
            })

            self.qsd.cmd('block-commit', {
                'job-id': 'commit-mid',
                'device': 'top',
                'top-node': 'mid',
                'base-node': 'node0',
                'auto-finalize': True,
                'auto-dismiss': False,
            })

            self._wait_for_blockjob('concluded')
            self.qsd.cmd('job-dismiss', {
                'id': 'commit-mid',
            })

            self.qsd.cmd('block-commit', {
                'job-id': 'commit-top',
                'device': 'top',
                'top-node': 'top',
                'base-node': 'node0',
                'auto-finalize': True,
                'auto-dismiss': False,
            })

            self._wait_for_blockjob('ready')
            self.qsd.cmd('job-complete', {
                'id': 'commit-top',
            })

            self._wait_for_blockjob('concluded')
            self.qsd.cmd('job-dismiss', {
                'id': 'commit-top',
            })

            self.qsd.cmd('blockdev-del', {
                'node-name': 'mid'
            })
            self.qsd.cmd('blockdev-del', {
                'node-name': 'top'
            })

        bench_thr.join()
        os.remove(mid)

if __name__ == '__main__':
    # Format must support raw backing files
    iotests.main(supported_fmts=['qcow', 'qcow2', 'qed'],
                 supported_protocols=['file'])
