nbd: Implement zero

Add zero(start, length) method sending NBD_CMD_WRITE_ZEROES command.

Update the example nbd tool to map data and zero ranges in the source
image and send zero areas efficiently.

Here is an example session uploading sparse images.

Creating 6g sparse source image:

$ virt-builder fedora-27 -o /var/tmp/fedora-27.img

Creating qcow2 destination image:

$ qemu-img create -f qcow2 /var/tmp/fedora-27.qcow2 6g

Serving image with qemu-nbd, supporting discard:

$ qemu-nbd \
      --socket /tmp/nbd.sock \
      --format qcow2 \
      --export-name=  \
      --persistent \
      --cache=none \
      --aio=native \
      --discard=unmap \
      /var/tmp/fedora-27.qcow2

Copying image to qcow2 image (best of 10 runs):

$ time PYTHONPATH=common examples/nbd-client upload /var/tmp/fedora-27.img /tmp/nbd.sock
[ 100.00% ] 6.00 GiB, 3.59 seconds, 1709.11 MiB/s

real	0m3.638s
user	0m0.072s
sys	0m0.688s

Same operation with qemu-img (best of 10 runs):

$ time qemu-img convert -n -p -f raw -O raw /var/tmp/fedora-27.img nbd:unix:/tmp/nbd.sock
    (100.00/100%)

real	0m5.435s
user	0m0.125s
sys	0m0.772s

Change-Id: I61462a6b8692948c573114800df31914c78c2e6c
Signed-off-by: Nir Soffer <nsoffer@redhat.com>

common/ovirt_imageio_common/nbd.py

@@ -59,6 +59,7 @@
 NBD_CMD_READ = 0
 NBD_CMD_WRITE = 1
 NBD_CMD_FLUSH = 3
+NBD_CMD_WRITE_ZEROES = 6
 
 # Error replies
 ERR_BASE = 2**31
@@ -141,6 +142,13 @@ def write(self, offset, data):
         self._send(data)
         self._receive_simple_reply(handle)
 
+    def zero(self, offset, length):
+        if self.transmission_flags & NBD_FLAG_SEND_WRITE_ZEROES == 0:
+            raise Error("Server does not support NBD_CMD_WRITE_ZEROES")
+        handle = next(self._counter)
+        self._send_command(NBD_CMD_WRITE_ZEROES, handle, offset, length)
+        self._receive_simple_reply(handle)
+
     def flush(self):
         # TODO: is this the best way to handle this?
         if self.transmission_flags & NBD_FLAG_SEND_FLUSH == 0:

common/test/nbd_test.py

@@ -25,12 +25,16 @@
 @contextmanager
 def nbd_server(image_path, image_format, sock, export_name="",
                read_only=False):
-    cmd = ["qemu-nbd",
-           "--socket", sock,
-           "--persistent",
-           "--cache", "none",
-           "--format", image_format,
-           "--export-name", export_name.encode("utf-8")]
+    cmd = [
+        "qemu-nbd",
+        "--socket", sock,
+        "--format", image_format,
+        "--export-name", export_name.encode("utf-8"),
+        "--persistent",
+        "--cache=none",
+        "--aio=native",
+        "--discard=unmap",
+    ]
 
     if read_only:
         cmd.append("--read-only")
@@ -128,3 +132,75 @@ def test_qcow2_write_read(tmpdir):
 
         with nbd.Client(sock) as c:
             assert c.read(offset, len(data)) == data
+
+
+@pytest.mark.parametrize("format", ["raw", "qcow2"])
+def test_zero(tmpdir, format):
+    size = 2 * 1024**2
+    offset = 1024**2
+    image = str(tmpdir.join("image"))
+    sock = str(tmpdir.join("sock"))
+    subprocess.check_call(
+        ["qemu-img", "create", "-f", format, image, str(size)])
+
+    with nbd_server(image, format, sock):
+        # Fill image with data
+        with nbd.Client(sock) as c:
+            c.write(0, b"x" * size)
+            c.flush()
+
+        # Zero a range
+        with nbd.Client(sock) as c:
+            c.zero(offset, 4096)
+            c.flush()
+
+        with nbd.Client(sock) as c:
+            assert c.read(offset, 4096) == b"\0" * 4096
+
+
+@pytest.mark.parametrize("format", ["raw", "qcow2"])
+def test_zero_max_block_size(tmpdir, format):
+    offset = 1024**2
+    image = str(tmpdir.join("image"))
+    sock = str(tmpdir.join("sock"))
+    subprocess.check_call(
+        ["qemu-img", "create", "-f", format, image, "1g"])
+
+    with nbd_server(image, format, sock):
+        # Fill range with data
+        with nbd.Client(sock) as c:
+            size = c.maximum_block_size
+            c.write(offset, b"x" * size)
+            c.flush()
+
+        # Zero range using maximum block size
+        with nbd.Client(sock) as c:
+            c.zero(offset, size)
+            c.flush()
+
+        with nbd.Client(sock) as c:
+            assert c.read(offset, size) == b"\0" * size
+
+
+@pytest.mark.parametrize("format", ["raw", "qcow2"])
+def test_zero_min_block_size(tmpdir, format):
+    offset = 1024**2
+    image = str(tmpdir.join("image"))
+    sock = str(tmpdir.join("sock"))
+    subprocess.check_call(
+        ["qemu-img", "create", "-f", format, image, "1g"])
+
+    with nbd_server(image, format, sock):
+        # Fill range with data
+        with nbd.Client(sock) as c:
+            size = c.minimum_block_size
+            c.write(offset, b"x" * size)
+            c.flush()
+
+        # Zero range using minimum block size
+        with nbd.Client(sock) as c:
+            c.zero(offset, size)
+            c.flush()
+
+        with nbd.Client(sock) as c:
+            assert c.read(offset, size) == b"\0" * size

examples/nbd-client

@@ -37,7 +37,9 @@ Upload source image to destination image via qemu-nbd:
 
 import argparse
 import io
+import json
 import os
+import subprocess
 import sys
 
 from ovirt_imageio_common import nbd
@@ -46,24 +48,55 @@ from ovirt_imageio_common import ui
 
 def upload(args):
     size = os.path.getsize(args.filename)
-    with io.open(args.filename, "rb") as f, \
-            nbd.Client(args.socket, args.export) as c, \
+    with io.open(args.filename, "rb") as src, \
+            nbd.Client(args.socket, args.export) as dst, \
             ui.ProgressBar(size) as pb:
-        if c.export_size < size:
+        if dst.export_size < size:
             raise Exception("Destination size {} is smaller than source file "
-                            "size {}".format(c.export_size, todo))
-        offset = 0
-        while offset < size:
-            chunk = f.read(min(size - offset, args.block_size))
-            if not chunk:
-                raise Exception("Unexpected end of file, expecting {} bytes"
-                                .format(todo))
-
-            c.write(offset, chunk)
-            offset += len(chunk)
-            pb.update(len(chunk))
-
-        c.flush()
+                            "size {}".format(dst.export_size, todo))
+        for zero, start, length in _map(args.filename):
+            if zero:
+                _zero_range(dst, start, length, pb)
+            else:
+                _copy_range(dst, src, start, length, pb, args.block_size)
+        dst.flush()
+
+
+def _map(path):
+    out = subprocess.check_output([
+        "qemu-img",
+        "map",
+        "--format", "raw",
+        "--output", "json",
+        path
+    ])
+    chunks = json.loads(out.decode("utf-8"))
+    for c in chunks:
+        yield c["zero"], c["start"], c["length"]
+
+
+def _zero_range(dst, start, length, pb):
+    while length:
+        step = min(dst.maximum_block_size, length)
+        dst.zero(start, step)
+        start += step
+        length -= step
+        pb.update(step)
+
+
+def _copy_range(dst, src, start, length, pb, block_size):
+    max_step = min(dst.maximum_block_size, block_size)
+    src.seek(start)
+    while length:
+        chunk = src.read(min(length, max_step))
+        if not chunk:
+            raise Exception("Unexpected end of file, expecting {} bytes"
+                            .format(length))
+        dst.write(start, chunk)
+        n = len(chunk)
+        start += n
+        length -= n
+        pb.update(n)
 
 
 def kib(s):
@@ -76,7 +109,7 @@ parser.add_argument(
     default="", help="export name")
 parser.add_argument(
     "-b", "--block-size",
-    default=1024**2,
+    default=8 * 1024**2,
     type=kib,
     help="block size in KiB")