allow DMA funcs to be called in quiesced state (#635)

Signed-off-by: Thanos Makatos <thanos.makatos@nutanix.com>
Reviewed-by: John Levon <john.levon@nutanix.com>

3 files changed, 168 insertions, 11 deletions

diff --git a/test/py/libvfio_user.py b/test/py/libvfio_user.py
index 60c5fbf..4b4a96a 100644
--- a/test/py/libvfio_user.py
+++ b/test/py/libvfio_user.py
@@ -672,9 +672,17 @@ def get_reply(sock, expect=0):
     return buf[16:]
 
 
-def msg(ctx, sock, cmd, payload, expect_reply_errno=0, fds=None, rsp=True,
-        expect_run_ctx_errno=None):
-    """Round trip a request and reply to the server."""
+def msg(ctx, sock, cmd, payload=bytearray(), expect_reply_errno=0, fds=None,
+        rsp=True, expect_run_ctx_errno=None):
+    """
+    Round trip a request and reply to the server. vfu_run_ctx will be
+    called once for the server to process the incoming message,
+    @expect_run_ctx_errrno checks the return value of vfu_run_ctx. If a
+    response is not expected then @rsp must be set to False, otherwise this
+    function will block indefinitely.
+    """
+    # FIXME if expect_run_ctx_errno == errno.EBUSY then shouldn't it implied
+    # that rsp == False?
     hdr = vfio_user_header(cmd, size=len(payload))
 
     if fds:
@@ -821,7 +829,7 @@ def dma_unregister(ctx, info):
 
 @vfu_dma_unregister_cb_t
 def __dma_unregister(ctx, info):
-    dma_unregister(ctx, info)
+    dma_unregister(ctx, copy.copy(info.contents))
 
 
 def quiesce_cb(ctx):
@@ -856,7 +864,7 @@ def vfu_setup_device_reset_cb(ctx, cb=_reset_cb):
 
 def prepare_ctx_for_dma(dma_register=__dma_register,
                         dma_unregister=__dma_unregister, quiesce=_quiesce_cb,
-                        reset=_reset_cb):
+                        reset=_reset_cb, migration_callbacks=False):
     ctx = vfu_create_ctx(flags=LIBVFIO_USER_FLAG_ATTACH_NB)
     assert ctx is not None
 
@@ -883,6 +891,10 @@ def prepare_ctx_for_dma(dma_register=__dma_register,
                            fd=f.fileno())
     assert ret == 0
 
+    if migration_callbacks:
+        ret = vfu_setup_device_migration_callbacks(ctx)
+        assert ret == 0
+
     ret = vfu_realize_ctx(ctx)
     assert ret == 0
 
@@ -949,10 +961,10 @@ def vfu_attach_ctx(ctx, expect=0):
 
 def vfu_run_ctx(ctx, expect_errno=None):
     ret = lib.vfu_run_ctx(ctx)
-    if expect_errno is not None:
+    if expect_errno:
         assert ret < 0 and expect_errno == c.get_errno(), \
-            "expected '%s' (%d), actual '%s' (%s)" % \
-                (os.strerror(expect_errno), expect_errno,
+            "ret=%s, expected '%s' (%d), actual '%s' (%s)" % \
+                (ret, os.strerror(expect_errno), expect_errno,
                  os.strerror(c.get_errno()), c.get_errno())
     return ret
 
@@ -1105,7 +1117,7 @@ def __migr_data_written_cb(ctx, count):
     return migr_data_written_cb(ctx, count)
 
 
-def vfu_setup_device_migration_callbacks(ctx, cbs=None, offset=0):
+def vfu_setup_device_migration_callbacks(ctx, cbs=None, offset=0x4000):
     assert ctx is not None
 
     if not cbs:
diff --git a/test/py/test_migration.py b/test/py/test_migration.py
index d70be5e..094a693 100644
--- a/test/py/test_migration.py
+++ b/test/py/test_migration.py
@@ -46,7 +46,7 @@ def setup_function(function):
                            flags=VFU_REGION_FLAG_RW)
     assert ret == 0
 
-    ret = vfu_setup_device_migration_callbacks(ctx, offset=0x4000)
+    ret = vfu_setup_device_migration_callbacks(ctx)
     assert ret == 0
 
     vfu_setup_device_quiesce_cb(ctx)
diff --git a/test/py/test_quiesce.py b/test/py/test_quiesce.py
index cf41e36..f231170 100644
--- a/test/py/test_quiesce.py
+++ b/test/py/test_quiesce.py
@@ -29,6 +29,7 @@
 
 from libvfio_user import *
 import errno
+from unittest import mock
 from unittest.mock import patch
 
 
@@ -37,7 +38,7 @@ ctx = None
 
 def setup_function(function):
     global ctx, sock
-    ctx = prepare_ctx_for_dma()
+    ctx = prepare_ctx_for_dma(migration_callbacks=True)
     assert ctx is not None
     sock = connect_client(ctx)
 
@@ -106,4 +107,148 @@ def test_device_quiesce_error_after_busy(mock_quiesce, mock_dma_register):
     assert c.get_errno() == errno.ENOENT
 
 
+# DMA map/unmap, migration device state transition, and reset callbacks
+# have the same function signature in Python
+def _side_effect(ctx, _):
+    count, sgs = vfu_addr_to_sg(ctx, 0x10000, 0x1000)
+    assert count == 1
+    sg = sgs[0]
+    assert sg.dma_addr == 0x10000 and sg.region == 0 \
+        and sg.length == 0x1000 and sg.offset == 0 and sg.writeable
+    iovec = iovec_t()
+    ret = vfu_map_sg(ctx, sg, iovec)
+    assert ret == 0, "%s" % c.get_errno()
+    assert iovec.iov_base != 0
+    assert iovec.iov_len == 0x1000
+    assert ret == 0
+    vfu_unmap_sg(ctx, sg, iovec)
+    return 0
+
+
+def _map_dma_region(ctx, sock, expect_run_ctx_errno=0):
+    rsp = expect_run_ctx_errno != errno.EBUSY
+    f = tempfile.TemporaryFile()
+    f.truncate(0x1000)
+    map_payload = vfio_user_dma_map(argsz=len(vfio_user_dma_map()),
+        flags=(VFIO_USER_F_DMA_REGION_READ | VFIO_USER_F_DMA_REGION_WRITE),
+        offset=0, addr=0x10000, size=0x1000)
+    msg(ctx, sock, VFIO_USER_DMA_MAP, map_payload, rsp=rsp,
+        expect_run_ctx_errno=expect_run_ctx_errno, fds=[f.fileno()])
+
+
+def _unmap_dma_region(ctx, sock, expect_run_ctx_errno=0):
+    rsp = expect_run_ctx_errno != errno.EBUSY
+    unmap_payload = vfio_user_dma_unmap(argsz=len(vfio_user_dma_unmap()),
+                                  addr=0x10000, size=0x1000)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, unmap_payload, rsp=rsp,
+        expect_run_ctx_errno=expect_run_ctx_errno)
+
+
+@patch('libvfio_user.dma_register', side_effect=_side_effect)
+@patch('libvfio_user.quiesce_cb')
+def test_allowed_funcs_in_quiesced_dma_register(mock_quiesce,
+                                                mock_dma_register):
+
+    global ctx, sock
+
+    # FIXME assert quiesce callback is called
+    _map_dma_region(ctx, sock)
+    # FIXME it's difficult to check that mock_dma_register has been called with
+    # the expected DMA info because we don't know the vaddr and the mapping
+    # (2nd and 3rd arguments of vfu_dma_info_t) as they're values returned from
+    # mmap(0) so they can't be predicted. Using mock.ANY in their place fails
+    # with "TypeError: cannot be converted to pointer". In any case this is
+    # tested by other unit tests.
+    mock_dma_register.assert_called_once_with(ctx, mock.ANY)
+
+
+@patch('libvfio_user.dma_unregister', side_effect=_side_effect)
+@patch('libvfio_user.quiesce_cb')
+def test_allowed_funcs_in_quiesced_dma_unregister(mock_quiesce,
+                                                  mock_dma_unregister):
+
+    global ctx, sock
+    _map_dma_region(ctx, sock)
+    _unmap_dma_region(ctx, sock)
+    mock_dma_unregister.assert_called_once_with(ctx, mock.ANY)
+
+
+@patch('libvfio_user.dma_register', side_effect=_side_effect)
+@patch('libvfio_user.quiesce_cb', side_effect=fail_with_errno(errno.EBUSY))
+def test_allowed_funcs_in_quiesced_dma_register_busy(mock_quiesce,
+                                                     mock_dma_register):
+
+    global ctx, sock
+    _map_dma_region(ctx, sock, errno.EBUSY)
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
+    mock_dma_register.assert_called_once_with(ctx, mock.ANY)
+
+
+@patch('libvfio_user.dma_unregister', side_effect=_side_effect)
+@patch('libvfio_user.quiesce_cb')
+def test_allowed_funcs_in_quiesced_dma_unregister_busy(mock_quiesce,
+                                                       mock_dma_unregister):
+
+    global ctx, sock
+    _map_dma_region(ctx, sock)
+    mock_quiesce.side_effect = fail_with_errno(errno.EBUSY)
+    _unmap_dma_region(ctx, sock, expect_run_ctx_errno=errno.EBUSY)
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
+    mock_dma_unregister.assert_called_once_with(ctx, mock.ANY)
+
+
+@patch('libvfio_user.migr_trans_cb', side_effect=_side_effect)
+@patch('libvfio_user.quiesce_cb')
+def test_allowed_funcs_in_quiesed_migration(mock_quiesce,
+                                            mock_trans):
+
+    global ctx, sock
+    _map_dma_region(ctx, sock)
+    data = VFIO_DEVICE_STATE_SAVING.to_bytes(c.sizeof(c.c_int), 'little')
+    write_region(ctx, sock, VFU_PCI_DEV_MIGR_REGION_IDX, offset=0,
+                 count=len(data), data=data)
+    mock_trans.assert_called_once_with(ctx, VFIO_DEVICE_STATE_SAVING)
+
+
+@patch('libvfio_user.migr_trans_cb', side_effect=_side_effect)
+@patch('libvfio_user.quiesce_cb')
+def test_allowed_funcs_in_quiesed_migration_busy(mock_quiesce,
+                                                 mock_trans):
+
+    global ctx, sock
+    _map_dma_region(ctx, sock)
+    mock_quiesce.side_effect = fail_with_errno(errno.EBUSY)
+    data = VFIO_DEVICE_STATE_STOP.to_bytes(c.sizeof(c.c_int), 'little')
+    write_region(ctx, sock, VFU_PCI_DEV_MIGR_REGION_IDX, offset=0,
+                 count=len(data), data=data, rsp=False,
+                 expect_run_ctx_errno=errno.EBUSY)
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
+    mock_trans.assert_called_once_with(ctx, VFIO_DEVICE_STATE_STOP)
+
+
+@patch('libvfio_user.reset_cb', side_effect=_side_effect)
+@patch('libvfio_user.quiesce_cb')
+def test_allowed_funcs_in_quiesced_reset(mock_quiesce, mock_reset):
+    global ctx, sock
+    _map_dma_region(ctx, sock)
+    msg(ctx, sock, VFIO_USER_DEVICE_RESET)
+    mock_reset.assert_called_once_with(ctx, VFU_RESET_DEVICE)
+
+
+@patch('libvfio_user.reset_cb', side_effect=_side_effect)
+@patch('libvfio_user.quiesce_cb')
+def test_allowed_funcs_in_quiesced_reset_busy(mock_quiesce, mock_reset):
+    global ctx, sock
+    _map_dma_region(ctx, sock)
+    mock_quiesce.side_effect = fail_with_errno(errno.EBUSY)
+    msg(ctx, sock, VFIO_USER_DEVICE_RESET, rsp=False,
+        expect_run_ctx_errno=errno.EBUSY)
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
+    mock_reset.assert_called_once_with(ctx, VFU_RESET_DEVICE)
+
+
 # ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #