Add Image.Extent() interface

Extent() return the next extent in the image, starting at the specified
offset, and limited by the specified length. An extent is a run of
clusters of with the same status.

Using Extent() we can iterate over the image and skip extents that reads
as zeros, instead of reading a zero buffer and detecting that the buffer
is full of zeros.

Benchmarking show 2 orders of magnitude speeded compared with zero
detection (10 TiB/s instead of 100 GiB/s)

    % go test -bench Benchmark0p/qcow2$/read
    Benchmark0p/qcow2/read-12      516     2503752 ns/op      107213.28 MB/s    1050537 B/op    39 allocs/op

    % go test -bench NextUnallocated
    BenchmarkNextUnallocated-12    116    10236186 ns/op    10489666.91 MB/s          0 B/op     0 allocs/op

Only qcow2 image has a real implementation. For other image formats we
treat all clusters as allocated, so Extent() always return one allocated
extent. For raw format we can implement Extents() later using SEEK_DATA
and SEEK_HOLE.

Extent() is more strict than ReadAt and fails when offset+length are
after the end of the image. We always know the length of the image so
requesting an extent after the end of the image is likely a bug in the
user code and failing fast may help to debug such issues.

Signed-off-by: Nir Soffer <nirsof@gmail.com>

Author: nirs

image/image.go

@@ -8,10 +8,28 @@ import (
 // Type must be a "Backing file format name string" that appears in QCOW2.
 type Type string
 
+// Extent describes a byte range in the image with the same allocation,
+// compression, or zero status. Extents are aligned to the underlying file
+// system block size (e.g. 4k), or the image format cluster size (e.g. 64k). One
+// extent can describe one or more file system blocks or image clusters.
+type Extent struct {
+	// Offset from start of the image in bytes.
+	Start int64 `json:"start"`
+	// Length of this extent in bytes.
+	Length int64 `json:"length"`
+	// Set if this extent is allocated.
+	Allocated bool `json:"allocated"`
+	// Set if this extent is read as zeros.
+	Zero bool `json:"zero"`
+	// Set if this extent is compressed.
+	Compressed bool `json:"compressed"`
+}
+
 // Image implements [io.ReaderAt] and [io.Closer].
 type Image interface {
 	io.ReaderAt
 	io.Closer
+	Extent(start, length int64) (Extent, error)
 	Type() Type
 	Size() int64 // -1 if unknown
 	Readable() error

image/qcow2/qcow2.go

@@ -776,6 +776,9 @@ func (img *Qcow2) getClusterMeta(off int64, cm *clusterMeta) error {
 	if cm.L2Entry.compressed() {
 		cm.Compressed = true
 	} else {
+		// When using extended L2 clusters this is always false. To find which sub
+		// cluster is allocated/zero we need to iterate over the allocation bitmap in
+		// the extended l2 cluster entry.
 		cm.Zero = standardClusterDescriptor(desc).allZero()
 	}
 
@@ -966,6 +969,110 @@ func readZero(p []byte, off int64, sz uint64) (int, error) {
 	return l, err
 }
 
+// clusterStatus returns an extent describing a single cluster. off must be aligned to
+// cluster size.
+func (img *Qcow2) clusterStatus(off int64) (image.Extent, error) {
+	var cm clusterMeta
+	if err := img.getClusterMeta(off, &cm); err != nil {
+		return image.Extent{}, err
+	}
+
+	if !cm.Allocated {
+		// If there is no backing file, or the cluster cannot be in the backing file,
+		// return an unallocated cluster.
+		if img.backingImage == nil || off >= img.backingImage.Size() {
+			// Unallocated cluster reads as zeros.
+			unallocated := image.Extent{Start: off, Length: int64(img.clusterSize), Zero: true}
+			return unallocated, nil
+		}
+
+		// Get the cluster from the backing file.
+		length := int64(img.clusterSize)
+		if off+length > img.backingImage.Size() {
+			length = img.backingImage.Size() - off
+		}
+		parent, err := img.backingImage.Extent(off, length)
+		if err != nil {
+			return parent, err
+		}
+		// The backing image may be a raw image not aligned to cluster size.
+		parent.Length = int64(img.clusterSize)
+		return parent, nil
+	}
+
+	// Cluster present in this image.
+	allocated := image.Extent{
+		Start:      off,
+		Length:     int64(img.clusterSize),
+		Allocated:  true,
+		Compressed: cm.Compressed,
+		Zero:       cm.Zero,
+	}
+	return allocated, nil
+}
+
+// Return true if extents have the same status.
+func sameStatus(a, b image.Extent) bool {
+	return a.Allocated == b.Allocated && a.Zero == b.Zero && a.Compressed == b.Compressed
+}
+
+// Extent returns the next extent starting at the specified offset. An extent
+// describes one or more clusters having the same status. The maximum length of
+// the returned extent is limited by the specified length. The minimum length of
+// the returned extent is length of one cluster.
+func (img *Qcow2) Extent(start, length int64) (image.Extent, error) {
+	// Default to zero length non-existent cluster.
+	var current image.Extent
+
+	if img.errUnreadable != nil {
+		return current, img.errUnreadable
+	}
+	if img.clusterSize == 0 {
+		return current, errors.New("cluster size cannot be 0")
+	}
+	if start+length > int64(img.Header.Size) {
+		return current, errors.New("length out of bounds")
+	}
+
+	// Compute the clusterStart of the first cluster to query. This may be behind start.
+	clusterStart := start / int64(img.clusterSize) * int64(img.clusterSize)
+
+	remaining := length
+	for remaining > 0 {
+		clusterStatus, err := img.clusterStatus(clusterStart)
+		if err != nil {
+			return current, err
+		}
+
+		// First cluster: if start is not aligned to cluster size, clip the start.
+		if clusterStatus.Start < start {
+			clusterStatus.Start = start
+			clusterStatus.Length -= start - clusterStatus.Start
+		}
+
+		// Last cluster: if start+length is not aligned to cluster size, clip the end.
+		if remaining < int64(img.clusterSize) {
+			clusterStatus.Length -= int64(img.clusterSize) - remaining
+		}
+
+		if current.Length == 0 {
+			// First cluster: copy status to current.
+			current = clusterStatus
+		} else if sameStatus(current, clusterStatus) {
+			// Cluster with same status: extend current.
+			current.Length += clusterStatus.Length
+		} else {
+			// Start of next extent
+			break
+		}
+
+		clusterStart += int64(img.clusterSize)
+		remaining -= clusterStatus.Length
+	}
+
+	return current, nil
+}
+
 // ReadAt implements [io.ReaderAt].
 func (img *Qcow2) ReadAt(p []byte, off int64) (n int, err error) {
 	if img.errUnreadable != nil {

image/raw/raw.go

@@ -1,6 +1,7 @@
 package raw
 
 import (
+	"errors"
 	"io"
 	"os"
 
@@ -14,6 +15,18 @@ type Raw struct {
 	io.ReaderAt `json:"-"`
 }
 
+// Extent returns an allocated extent starting at the specified offset with
+// specified length. It is used when the speicfic image type does not implement
+// Extent(). The implementation is correct but inefficient. Fails if image size
+// is unknown.
+func (img *Raw) Extent(start, length int64) (image.Extent, error) {
+	if start+length > img.Size() {
+		return image.Extent{}, errors.New("length out of bounds")
+	}
+	// TODO: Implement using SEEK_HOLE/SEEK_DATA when supported by the file system.
+	return image.Extent{Start: start, Length: length, Allocated: true}, nil
+}
+
 func (img *Raw) Close() error {
 	if closer, ok := img.ReaderAt.(io.Closer); ok {
 		return closer.Close()

image/stub/stub.go

@@ -25,6 +25,10 @@ type Stub struct {
 	t image.Type
 }
 
+func (img *Stub) Extent(start, length int64) (image.Extent, error) {
+	return image.Extent{}, fmt.Errorf("unimplemented type: %q", img.t)
+}
+
 func (img *Stub) ReadAt([]byte, int64) (int, error) {
 	return 0, img.Readable()
 }

qcow2reader_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/lima-vm/go-qcow2reader"
 	"github.com/lima-vm/go-qcow2reader/convert"
+	"github.com/lima-vm/go-qcow2reader/image"
 	"github.com/lima-vm/go-qcow2reader/test/qemuimg"
 )
 
@@ -19,6 +20,186 @@ const (
 	GiB = int64(1) << 30
 )
 
+func TestExtentsUnallocated(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "image")
+	if err := qemuimg.Create(path, qemuimg.FormatQcow2, 4*GiB, "", ""); err != nil {
+		t.Fatal(err)
+	}
+	f, err := os.Open(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+	img, err := qcow2reader.Open(f)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer img.Close()
+
+	t.Run("entire image", func(t *testing.T) {
+		actual, err := img.Extent(0, img.Size())
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := image.Extent{Start: 0, Length: img.Size(), Zero: true}
+		if actual != expected {
+			t.Fatalf("expected %+v, got %+v", expected, actual)
+		}
+	})
+	t.Run("same result", func(t *testing.T) {
+		r1, err := img.Extent(0, img.Size())
+		if err != nil {
+			t.Fatal(err)
+		}
+		r2, err := img.Extent(0, img.Size())
+		if err != nil {
+			t.Fatal(err)
+		}
+		if r1 != r2 {
+			t.Fatalf("expected %+v, got %+v", r1, r2)
+		}
+	})
+	t.Run("all segments", func(t *testing.T) {
+		for i := int64(0); i < img.Size(); i += 32 * MiB {
+			segment, err := img.Extent(i, 32*MiB)
+			if err != nil {
+				t.Fatal(err)
+			}
+			expected := image.Extent{Start: i, Length: 32 * MiB, Zero: true}
+			if segment != expected {
+				t.Fatalf("expected %+v, got %+v", expected, segment)
+			}
+		}
+	})
+	t.Run("start unaligned", func(t *testing.T) {
+		start := 32*MiB + 42
+		length := 32 * MiB
+		actual, err := img.Extent(start, length)
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := image.Extent{Start: start, Length: length, Zero: true}
+		if actual != expected {
+			t.Fatalf("expected %+v, got %+v", expected, actual)
+		}
+	})
+	t.Run("length unaligned", func(t *testing.T) {
+		start := 32 * MiB
+		length := 32*MiB - 42
+		actual, err := img.Extent(start, length)
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := image.Extent{Start: start, Length: length, Zero: true}
+		if actual != expected {
+			t.Fatalf("expected %+v, got %+v", expected, actual)
+		}
+	})
+	t.Run("start and length unaligned", func(t *testing.T) {
+		start := 32*MiB + 42
+		length := 32*MiB - 42
+		actual, err := img.Extent(start, length)
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := image.Extent{Start: start, Length: length, Zero: true}
+		if actual != expected {
+			t.Fatalf("expected %+v, got %+v", expected, actual)
+		}
+	})
+	t.Run("length after end of image", func(t *testing.T) {
+		start := img.Size() - 31*MiB
+		actual, err := img.Extent(start, 32*MiB)
+		if err == nil {
+			t.Fatal("out of bounds request did not fail")
+		}
+		var expected image.Extent
+		if actual != expected {
+			t.Fatalf("expected %+v, got %+v", expected, actual)
+		}
+	})
+	t.Run("start after end of image", func(t *testing.T) {
+		start := img.Size() + 1*MiB
+		actual, err := img.Extent(start, 32*MiB)
+		if err == nil {
+			t.Fatal("out of bounds request did not fail")
+		}
+		var expected image.Extent
+		if actual != expected {
+			t.Fatalf("expected %+v, got %+v", expected, actual)
+		}
+	})
+}
+
+func TestExtentsRaw(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "disk.img")
+	size := 4 * GiB
+	f, err := os.Create(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+	if err := f.Truncate(size); err != nil {
+		t.Fatal(err)
+	}
+	img, err := qcow2reader.Open(f)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer img.Close()
+
+	t.Run("entire image", func(t *testing.T) {
+		actual, err := img.Extent(0, img.Size())
+		if err != nil {
+			t.Fatal(err)
+		}
+		// Currently we always report raw images as fully allocated.
+		expected := image.Extent{Start: 0, Length: img.Size(), Allocated: true}
+		if actual != expected {
+			t.Fatalf("expected %+v, got %+v", expected, actual)
+		}
+	})
+	t.Run("length after end of image", func(t *testing.T) {
+		start := img.Size() - 31*MiB
+		actual, err := img.Extent(start, 32*MiB)
+		if err == nil {
+			t.Fatal("out of bounds request did not fail")
+		}
+		var expected image.Extent
+		if actual != expected {
+			t.Fatalf("expected %+v, got %+v", expected, actual)
+		}
+	})
+}
+
+func BenchmarkExtentsUnallocated(b *testing.B) {
+	path := filepath.Join(b.TempDir(), "image")
+	if err := qemuimg.Create(path, qemuimg.FormatQcow2, 100*GiB, "", ""); err != nil {
+		b.Fatal(err)
+	}
+	f, err := os.Open(path)
+	if err != nil {
+		b.Fatal(err)
+	}
+	img, err := qcow2reader.Open(f)
+	if err != nil {
+		b.Fatal(err)
+	}
+	expected := image.Extent{Start: 0, Length: img.Size(), Zero: true}
+	resetBenchmark(b, img.Size())
+	for i := 0; i < b.N; i++ {
+		b.StartTimer()
+		actual, err := img.Extent(0, img.Size())
+		b.StopTimer()
+		if err != nil {
+			b.Fatal(err)
+		}
+		if actual != expected {
+			b.Fatalf("expected %+v, got %+v", expected, actual)
+		}
+	}
+}
+
 // Benchmark completely empty sparse image (0% utilization).  This is the best
 // case when we don't have to read any cluster from storage.
 func Benchmark0p(b *testing.B) {