Merge pull request #31 from nirs/cache-l2-tables

Cache l2 tables

Author: AkihiroSuda

image/qcow2/qcow2.go

@@ -14,6 +14,7 @@ import (
 	"github.com/lima-vm/go-qcow2reader/align"
 	"github.com/lima-vm/go-qcow2reader/image"
 	"github.com/lima-vm/go-qcow2reader/log"
+	"github.com/lima-vm/go-qcow2reader/lru"
 )
 
 const Type = "qcow2"
@@ -536,20 +537,25 @@ type Qcow2 struct {
 	errUnreadable       error
 	clusterSize         int
 	l1Table             []l1TableEntry
+	l2TableCache        *lru.Cache[l1TableEntry, []l2TableEntry]
 	decompressor        Decompressor
 	BackingFile         string     `json:"backing_file"`
 	BackingFileFullPath string     `json:"backing_file_full_path"`
 	BackingFileFormat   image.Type `json:"backing_file_format"`
 	backingImage        image.Image
 }
 
+// With the default cluster size (64 Kib) this uses 1 MiB and cover 8 GiB image.
+const maxL2Tables = 16
+
 // Open opens an qcow2 image.
 //
 // To open an image with backing files, ra must implement [Namer],
 // and openWithType must be non-nil.
 func Open(ra io.ReaderAt, openWithType image.OpenWithType) (*Qcow2, error) {
 	img := &Qcow2{
-		ra: ra,
+		ra:           ra,
+		l2TableCache: lru.New[l1TableEntry, []l2TableEntry](maxL2Tables),
 	}
 	r := io.NewSectionReader(ra, 0, -1)
 	var err error
@@ -680,6 +686,19 @@ func (img *Qcow2) extendedL2() bool {
 	return img.Header.HeaderFieldsV3 != nil && img.Header.HeaderFieldsV3.IncompatibleFeatures&(1<<IncompatibleFeaturesExtendedL2EntriesBit) != 0
 }
 
+func (img *Qcow2) getL2Table(l1Entry l1TableEntry) ([]l2TableEntry, error) {
+	l2Table, ok := img.l2TableCache.Get(l1Entry)
+	if !ok {
+		var err error
+		l2Table, err = readL2Table(img.ra, l1Entry.l2Offset(), img.clusterSize)
+		if err != nil {
+			return nil, err
+		}
+		img.l2TableCache.Add(l1Entry, l2Table)
+	}
+	return l2Table, nil
+}
+
 // readAtAligned requires that off and off+len(p)-1 belong to the same cluster.
 func (img *Qcow2) readAtAligned(p []byte, off int64) (int, error) {
 	l2Entries := img.clusterSize / 8
@@ -712,7 +731,7 @@ func (img *Qcow2) readAtAligned(p []byte, off int64) (int, error) {
 		extL2Entry = &extL2Table[l2Index]
 		l2Entry = extL2Entry.L2TableEntry
 	} else {
-		l2Table, err := readL2Table(img.ra, l2TableOffset, img.clusterSize)
+		l2Table, err := img.getL2Table(l1Entry)
 		if err != nil {
 			return 0, fmt.Errorf("failed to read L2 table for L1 entry %v (index %d): %w", l1Entry, l1Index, err)
 		}

lru/lru.go

@@ -0,0 +1,70 @@
+package lru
+
+import (
+	"container/list"
+	"sync"
+)
+
+// Cache keeps recently used values. Safe for concurrent use by multiple
+// goroutines.
+type Cache[K comparable, V any] struct {
+	mutex        sync.Mutex
+	entries      map[K]*list.Element
+	recentlyUsed *list.List
+	capacity     int
+}
+
+type cacheEntry[K comparable, V any] struct {
+	Key   K
+	Value V
+}
+
+// New returns a new empty cache that can hold up to capacity items.
+func New[K comparable, V any](capacity int) *Cache[K, V] {
+	return &Cache[K, V]{
+		entries:      make(map[K]*list.Element),
+		recentlyUsed: list.New(),
+		capacity:     capacity,
+	}
+}
+
+// Get returns the value stored in the cache for a key, or zero value if no
+// value is present. The ok result indicates whether value was found in the
+// cache.
+func (c *Cache[K, V]) Get(key K) (V, bool) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	if elem, ok := c.entries[key]; ok {
+		c.recentlyUsed.MoveToFront(elem)
+		entry := elem.Value.(*cacheEntry[K, V])
+		return entry.Value, true
+	}
+
+	var missing V
+	return missing, false
+}
+
+// Add adds key and value to the cache. If the cache is full, the oldest entry
+// is removed. If key is already in the cache, value replaces the cached value.
+func (c *Cache[K, V]) Add(key K, value V) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	if elem, ok := c.entries[key]; ok {
+		c.recentlyUsed.MoveToFront(elem)
+		entry := elem.Value.(*cacheEntry[K, V])
+		entry.Value = value
+		return
+	}
+
+	if len(c.entries) >= c.capacity {
+		oldest := c.recentlyUsed.Back()
+		c.recentlyUsed.Remove(oldest)
+		entry := oldest.Value.(*cacheEntry[K, V])
+		delete(c.entries, entry.Key)
+	}
+
+	entry := &cacheEntry[K, V]{Key: key, Value: value}
+	c.entries[key] = c.recentlyUsed.PushFront(entry)
+}

lru/lru_test.go

@@ -0,0 +1,61 @@
+package lru
+
+import (
+	"testing"
+)
+
+type item struct {
+	key   int
+	value string
+}
+
+func TestCache(t *testing.T) {
+	items := []item{{0, "0"}, {1, "1"}, {2, "2"}, {3, "3"}, {4, "4"}}
+	cache := New[int, string](5)
+
+	// Cache is empty.
+	for _, i := range items {
+		if _, ok := cache.Get(i.key); ok {
+			t.Errorf("key %d in cache", i.key)
+		}
+	}
+
+	// Add all items to cache.
+	for _, i := range items {
+		cache.Add(i.key, i.value)
+	}
+
+	// Verify that all items are cached.
+	for _, i := range items {
+		if value, ok := cache.Get(i.key); !ok {
+			t.Errorf("cached value for %d missing", i.key)
+		} else if value != i.value {
+			t.Errorf("expected %q, got %q", value, i.value)
+		}
+	}
+
+	// Adding next item will remove the least used item (0).
+	cache.Add(5, "5")
+
+	// New item in cache.
+	if value, ok := cache.Get(5); !ok {
+		t.Errorf("cached value for 5 missing")
+	} else if value != "5" {
+		t.Errorf("expected \"5\", got %q", value)
+	}
+
+	// Removed item not in cache.
+	if _, ok := cache.Get(0); ok {
+		t.Error("key 0 in cache")
+	}
+
+	// Rest of items not affected.
+	for _, i := range items[1:] {
+		if value, ok := cache.Get(i.key); !ok {
+			t.Errorf("cached value for %d missing", i.key)
+		} else if value != i.value {
+			t.Errorf("expected %q, got %q", value, i.value)
+		}
+	}
+
+}

qcow2reader_test.go

@@ -0,0 +1,149 @@
+package qcow2reader
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"io"
+	"math/rand"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"testing"
+
+	"github.com/lima-vm/go-qcow2reader/image"
+	"github.com/lima-vm/go-qcow2reader/image/qcow2"
+)
+
+const (
+	MiB = int64(1) << 20
+	GiB = int64(1) << 30
+
+	CompressionTypeNone = qcow2.CompressionType(255)
+)
+
+func BenchmarkRead(b *testing.B) {
+	const size = 256 * MiB
+	base := filepath.Join(b.TempDir(), "image")
+	if err := createTestImage(base, size); err != nil {
+		b.Fatal(err)
+	}
+	b.Run("qcow2", func(b *testing.B) {
+		img := base + ".qocw2"
+		if err := qemuImgConvert(base, img, qcow2.Type, CompressionTypeNone); err != nil {
+			b.Fatal(err)
+		}
+		resetBenchmark(b, size)
+		for i := 0; i < b.N; i++ {
+			benchmarkRead(b, img)
+		}
+	})
+	b.Run("qcow2 zlib", func(b *testing.B) {
+		img := base + ".zlib.qcow2"
+		if err := qemuImgConvert(base, img, qcow2.Type, qcow2.CompressionTypeZlib); err != nil {
+			b.Fatal(err)
+		}
+		resetBenchmark(b, size)
+		for i := 0; i < b.N; i++ {
+			benchmarkRead(b, img)
+		}
+	})
+	// TODO: qcow2 zstd (not supported yet)
+}
+
+func benchmarkRead(b *testing.B, filename string) {
+	b.StartTimer()
+
+	f, err := os.Open(filename)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer f.Close()
+	img, err := Open(f)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer img.Close()
+	buf := make([]byte, 1*MiB)
+	reader := io.NewSectionReader(img, 0, img.Size())
+	n, err := io.CopyBuffer(io.Discard, reader, buf)
+
+	b.StopTimer()
+
+	if err != nil {
+		b.Fatal(err)
+	}
+	if n != img.Size() {
+		b.Fatalf("Expected %d bytes, read %d bytes", img.Size(), n)
+	}
+}
+
+func resetBenchmark(b *testing.B, size int64) {
+	b.StopTimer()
+	b.ResetTimer()
+	b.SetBytes(size)
+	b.ReportAllocs()
+}
+
+// createTestImage creates a 50% allocated raw image with fake data that
+// compresses like real image data.
+func createTestImage(filename string, size int64) error {
+	const chunkSize = 4 * MiB
+	file, err := os.Create(filename)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+	if err := file.Truncate(size); err != nil {
+		return err
+	}
+	reader := &Generator{}
+	for offset := int64(0); offset < size; offset += 2 * chunkSize {
+		_, err := file.Seek(offset, io.SeekStart)
+		if err != nil {
+			return err
+		}
+		chunk := io.LimitReader(reader, chunkSize)
+		if n, err := io.Copy(file, chunk); err != nil {
+			return err
+		} else if n != chunkSize {
+			return fmt.Errorf("expected %d bytes, wrote %d bytes", chunkSize, n)
+		}
+	}
+	return file.Close()
+}
+
+// Generator generates fake data that compresses like a real image data (30%).
+type Generator struct{}
+
+func (g *Generator) Read(b []byte) (int, error) {
+	for i := 0; i < len(b); i++ {
+		b[i] = byte(i & 0xff)
+	}
+	rand.Shuffle(len(b)/8*5, func(i, j int) {
+		b[i], b[j] = b[j], b[i]
+	})
+	return len(b), nil
+}
+
+func qemuImgConvert(src, dst string, dstFormat image.Type, compressionType qcow2.CompressionType) error {
+	args := []string{"convert", "-O", string(dstFormat)}
+	if compressionType != CompressionTypeNone {
+		args = append(args, "-c", "-o", "compression_type="+compressionType.String())
+	}
+	args = append(args, src, dst)
+	cmd := exec.Command("qemu-img", args...)
+
+	var stderr bytes.Buffer
+	cmd.Stderr = &stderr
+
+	if err := cmd.Run(); err != nil {
+		// Return qemu-img stderr instead of the unhelpful default error (exited
+		// with status 1).
+		if _, ok := err.(*exec.ExitError); ok {
+			return errors.New(stderr.String())
+		}
+		return err
+	}
+	return nil
+}