sow-tools/internal/nwsync/memory_test.go

package nwsync

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"runtime"
	"runtime/debug"
	"testing"
	"time"

	"github.com/klauspost/compress/zstd"

	"git.westgate.pw/ShadowsOverWestgate/sow-tools/internal/erf"
)

// TestSingleThreadedEncoderMatchesDefault pins the claim the blob encoder's
// concurrency setting rests on: it saves memory only, and a published blob is
// the same bytes either way.
func TestSingleThreadedEncoderMatchesDefault(t *testing.T) {
	standard, err := zstd.NewWriter(nil)
	if err != nil {
		t.Fatal(err)
	}
	defer standard.Close()

	body := make([]byte, 4<<20)
	random := rand.New(rand.NewSource(1))
	random.Read(body[:len(body)/2])
	for _, size := range []int{0, 1, 4 << 10, len(body)} {
		if !bytes.Equal(blobEncoder.EncodeAll(body[:size], nil), standard.EncodeAll(body[:size], nil)) {
			t.Fatalf("%d bytes compress differently at concurrency 1", size)
		}
	}
}

// resourceSize is one payload in the memory fixtures. Real haks hold a few MB
// per resource, and peak memory is meant to track that, not the archive.
const resourceSize = 1 << 20

// writeStreamedHak builds a hak of count resources without ever holding the
// archive in memory, so the fixture itself does not decide the measurement.
// Payloads are distinct, so no blob is deduplicated away.
func writeStreamedHak(t *testing.T, path string, count int) {
	t.Helper()
	payload := filepath.Join(t.TempDir(), "payload.bin")
	body := make([]byte, resourceSize)
	for index := range body {
		body[index] = byte(index)
	}

	resources := make([]erf.Resource, 0, count)
	for index := range count {
		// A distinct first byte per resource is enough to give every payload
		// its own sha1 while still streaming from one file per resource.
		unique := filepath.Join(filepath.Dir(payload), fmt.Sprintf("p%d.bin", index))
		body[0] = byte(index)
		body[1] = byte(index >> 8)
		if err := os.WriteFile(unique, body, 0o644); err != nil {
			t.Fatal(err)
		}
		resources = append(resources, erf.Resource{
			Name:       fmt.Sprintf("res%05d", index),
			Type:       restype(t, "tga"),
			SourcePath: unique,
			Size:       resourceSize,
		})
	}

	file, err := os.Create(path)
	if err != nil {
		t.Fatal(err)
	}
	defer file.Close()
	if err := erf.Write(file, erf.New("HAK", resources)); err != nil {
		t.Fatalf("write hak: %v", err)
	}
}

// peakHeapDuring runs work while sampling the heap, and returns the largest
// live heap it saw.
func peakHeapDuring(work func()) uint64 {
	runtime.GC()
	done := make(chan struct{})
	peak := make(chan uint64, 1)
	go func() {
		var highest uint64
		var stats runtime.MemStats
		for {
			select {
			case <-done:
				peak <- highest
				return
			default:
			}
			runtime.ReadMemStats(&stats)
			if stats.HeapAlloc > highest {
				highest = stats.HeapAlloc
			}
			time.Sleep(time.Millisecond)
		}
	}()
	work()
	close(done)
	return <-peak
}

// TestEmitPeakMemoryDoesNotScaleWithArtifactSize is the regression check for
// the OOM kills on large haks: emit used to hold the whole archive (twice), so
// a 2 GB hak needed about 10 GB. Emitting an archive 8× bigger must not cost
// meaningfully more memory.
func TestEmitPeakMemoryDoesNotScaleWithArtifactSize(t *testing.T) {
	if testing.Short() {
		t.Skip("writes a 64 MB fixture")
	}
	// A lazy GC lets garbage pile up in proportion to the live heap, which
	// hides the thing under test. Collecting eagerly makes the sampled heap
	// track what emit actually holds.
	defer debug.SetGCPercent(debug.SetGCPercent(10))

	measure := func(count int) uint64 {
		dir := t.TempDir()
		hak := filepath.Join(dir, "big.hak")
		writeStreamedHak(t, hak, count)
		// The key is computed outside the measurement: the test helper reads
		// the whole file to hash it, which emit itself no longer does.
		options := EmitOptions{
			ArtifactKey:  artifactKey(t, hak),
			ArtifactPath: hak,
			As:           filepath.Base(hak),
			OutDir:       filepath.Join(dir, "out"),
		}
		return peakHeapDuring(func() {
			if _, err := Emit(options); err != nil {
				t.Fatalf("emit %d resources: %v", count, err)
			}
		})
	}

	small := measure(8)  // 8 MB
	large := measure(64) // 64 MB
	const slack = 24 << 20

	t.Logf("peak heap: 8 MB hak %d bytes, 64 MB hak %d bytes", small, large)
	if large > small+slack {
		t.Fatalf("peak heap scaled with artifact size: 8 MB hak peaked at %d bytes, 64 MB hak at %d", small, large)
	}
}