fix(nwsync): stream emit so peak memory tracks the largest resource (#76) #78

Merged
archvillainette merged 2 commits from fix/nwsync-emit-streaming into main 2026-07-31 11:16:23 +00:00
5 changed files with 285 additions and 92 deletions
Showing only changes of commit 59b233e3db - Show all commits
+91 -45
View File
@@ -311,63 +311,109 @@ func Write(w io.Writer, archive Archive) error {
return nil return nil
} }
// IndexEntry locates one resource inside an archive without holding its
// payload. Streaming callers read one payload at a time from these, so peak
// memory tracks the largest resource instead of the whole archive.
type IndexEntry struct {
Name string
Type uint16
Offset int64
Size int64
}
// Index is the header plus the resource table of an ERF: everything except the
// payloads.
type Index struct {
FileType string
Version string
Entries []IndexEntry
}
// ReadIndex parses the tables of an ERF of the given size, reading only the
// header, the key list and the resource list.
func ReadIndex(r io.ReaderAt, size int64) (Index, error) {
if size < headerSize {
return Index{}, fmt.Errorf("erf file too small: %d bytes", size)
}
var hdr header
if err := binary.Read(io.NewSectionReader(r, 0, headerSize), binary.LittleEndian, &hdr); err != nil {
return Index{}, fmt.Errorf("decode erf header: %w", err)
}
if int64(hdr.KeyListOffset)+int64(hdr.EntryCount)*24 > size {
return Index{}, fmt.Errorf("erf key list exceeds file bounds")
}
keys := make([]keyEntry, hdr.EntryCount)
keyReader := io.NewSectionReader(r, int64(hdr.KeyListOffset), int64(hdr.EntryCount)*24)
if err := binary.Read(keyReader, binary.LittleEndian, &keys); err != nil {
return Index{}, fmt.Errorf("decode key list: %w", err)
}
if int64(hdr.ResourceListOffset)+int64(hdr.EntryCount)*8 > size {
return Index{}, fmt.Errorf("erf resource list exceeds file bounds")
}
entries := make([]resourceEntry, hdr.EntryCount)
entryReader := io.NewSectionReader(r, int64(hdr.ResourceListOffset), int64(hdr.EntryCount)*8)
if err := binary.Read(entryReader, binary.LittleEndian, &entries); err != nil {
return Index{}, fmt.Errorf("decode resource list: %w", err)
}
index := Index{
FileType: string(hdr.FileType[:]),
Version: string(hdr.Version[:]),
Entries: make([]IndexEntry, 0, hdr.EntryCount),
}
for position, key := range keys {
entry := entries[position]
if int64(entry.Offset)+int64(entry.Size) > size {
return Index{}, fmt.Errorf("resource %d exceeds file bounds", position)
}
index.Entries = append(index.Entries, IndexEntry{
Name: string(bytes.TrimRight(key.ResRef[:], "\x00")),
Type: key.ResourceType,
Offset: int64(entry.Offset),
Size: int64(entry.Size),
})
}
return index, nil
}
// ReadPayload returns one resource's bytes.
func ReadPayload(r io.ReaderAt, entry IndexEntry) ([]byte, error) {
payload := make([]byte, entry.Size)
if _, err := r.ReadAt(payload, entry.Offset); err != nil {
return nil, fmt.Errorf("read resource %q: %w", entry.Name, err)
}
return payload, nil
}
// Read materialises a whole archive. Payloads are subslices of the buffer the
// archive was read into, so nothing is copied twice; callers that only need one
// resource at a time should use ReadIndex instead.
func Read(r io.Reader) (Archive, error) { func Read(r io.Reader) (Archive, error) {
data, err := io.ReadAll(r) data, err := io.ReadAll(r)
if err != nil { if err != nil {
return Archive{}, fmt.Errorf("read erf: %w", err) return Archive{}, fmt.Errorf("read erf: %w", err)
} }
if len(data) < headerSize { index, err := ReadIndex(bytes.NewReader(data), int64(len(data)))
return Archive{}, fmt.Errorf("erf file too small: %d bytes", len(data)) if err != nil {
return Archive{}, err
} }
var hdr header resources := make([]Resource, 0, len(index.Entries))
if err := binary.Read(bytes.NewReader(data[:headerSize]), binary.LittleEndian, &hdr); err != nil { for _, entry := range index.Entries {
return Archive{}, fmt.Errorf("decode erf header: %w", err)
}
keyStart := int(hdr.KeyListOffset)
keyEnd := keyStart + int(hdr.EntryCount)*24
if keyEnd > len(data) {
return Archive{}, fmt.Errorf("erf key list exceeds file bounds")
}
keys := make([]keyEntry, hdr.EntryCount)
if err := binary.Read(bytes.NewReader(data[keyStart:keyEnd]), binary.LittleEndian, &keys); err != nil {
return Archive{}, fmt.Errorf("decode key list: %w", err)
}
resourceStart := int(hdr.ResourceListOffset)
resourceEnd := resourceStart + int(hdr.EntryCount)*8
if resourceEnd > len(data) {
return Archive{}, fmt.Errorf("erf resource list exceeds file bounds")
}
entries := make([]resourceEntry, hdr.EntryCount)
if err := binary.Read(bytes.NewReader(data[resourceStart:resourceEnd]), binary.LittleEndian, &entries); err != nil {
return Archive{}, fmt.Errorf("decode resource list: %w", err)
}
resources := make([]Resource, 0, hdr.EntryCount)
for index, key := range keys {
entry := entries[index]
start := int(entry.Offset)
end := start + int(entry.Size)
if end > len(data) {
return Archive{}, fmt.Errorf("resource %d exceeds file bounds", index)
}
resref := string(bytes.TrimRight(key.ResRef[:], "\x00"))
payload := make([]byte, entry.Size)
copy(payload, data[start:end])
resources = append(resources, Resource{ resources = append(resources, Resource{
Name: resref, Name: entry.Name,
Type: key.ResourceType, Type: entry.Type,
Data: payload, Data: data[entry.Offset : entry.Offset+entry.Size],
Size: int64(entry.Size), Size: entry.Size,
}) })
} }
return Archive{ return Archive{
FileType: string(hdr.FileType[:]), FileType: index.FileType,
Version: string(hdr.Version[:]), Version: index.Version,
Resources: resources, Resources: resources,
}, nil }, nil
} }
+6 -2
View File
@@ -22,9 +22,13 @@ const (
blobHeaderBytes = 24 blobHeaderBytes = 24
) )
// EncodeAll/DecodeAll are single-threaded per call, so the default pool of one
// encoder per CPU only buys idle memory: each holds a window-sized history, so
// on a 24-core runner that is ~200 MB of live heap doing nothing. Concurrency 1
// produces byte-identical output.
var ( var (
blobEncoder, _ = zstd.NewWriter(nil) blobEncoder, _ = zstd.NewWriter(nil, zstd.WithEncoderConcurrency(1))
blobDecoder, _ = zstd.NewReader(nil) blobDecoder, _ = zstd.NewReader(nil, zstd.WithDecoderConcurrency(1))
) )
// compressBlob wraps data in NWCompressedBuffer framing. // compressBlob wraps data in NWCompressedBuffer framing.
+54 -42
View File
@@ -1,10 +1,10 @@
package nwsync package nwsync
import ( import (
"bytes"
"context" "context"
"crypto/sha1" "crypto/sha1"
"fmt" "fmt"
"io"
"os" "os"
"path" "path"
"path/filepath" "path/filepath"
@@ -79,7 +79,12 @@ type EmitOptions struct {
// real blobs in the zone and no index, which is unambiguous. Blob names are // real blobs in the zone and no index, which is unambiguous. Blob names are
// content hashes, so re-running skips whatever already landed. // content hashes, so re-running skips whatever already landed.
func Emit(options EmitOptions) (EmitResult, error) { func Emit(options EmitOptions) (EmitResult, error) {
artifact, err := os.ReadFile(options.ArtifactPath) artifact, err := os.Open(options.ArtifactPath)
if err != nil {
return EmitResult{}, fmt.Errorf("read artifact: %w", err)
}
defer artifact.Close()
info, err := artifact.Stat()
if err != nil { if err != nil {
return EmitResult{}, fmt.Errorf("read artifact: %w", err) return EmitResult{}, fmt.Errorf("read artifact: %w", err)
} }
@@ -98,7 +103,7 @@ func Emit(options EmitOptions) (EmitResult, error) {
return EmitResult{}, err return EmitResult{}, err
} }
resources, err := readArtifact(options.ArtifactPath, artifact, name) index, err := readArtifactIndex(options.ArtifactPath, artifact, info.Size(), name)
if err != nil { if err != nil {
return EmitResult{}, err return EmitResult{}, err
} }
@@ -108,7 +113,7 @@ func Emit(options EmitOptions) (EmitResult, error) {
return EmitResult{}, err return EmitResult{}, err
} }
entries, blobs, onDiskBytes, err := emitResources(resources, target) entries, blobs, onDiskBytes, err := emitResources(artifact, index, target)
if err != nil { if err != nil {
return EmitResult{}, err return EmitResult{}, err
} }
@@ -137,60 +142,64 @@ func openSink(outDir string, injected sink) (sink, error) {
return newZoneSink(context.Background(), os.Getenv) return newZoneSink(context.Background(), os.Getenv)
} }
// readArtifact returns the resources of an ERF/HAK/MOD, or the single resource // readArtifactIndex locates the resources of an ERF/HAK/MOD, or the single
// a loose file represents. Upstream's resman does the same dispatch on the // resource a loose file represents, without reading any payload. Upstream's
// file's first three bytes. name is the artifact's published name, which for a // resman does the same dispatch on the file's first three bytes. name is the
// loose file is also its resref. // artifact's published name, which for a loose file is also its resref.
func readArtifact(path string, data []byte, name string) ([]erf.Resource, error) { func readArtifactIndex(path string, artifact io.ReaderAt, size int64, name string) ([]erf.IndexEntry, error) {
if len(data) >= 3 { magic := make([]byte, 3)
switch string(data[:3]) { if size >= 3 {
case "ERF", "HAK": if _, err := artifact.ReadAt(magic, 0); err != nil {
archive, err := erf.Read(bytes.NewReader(data)) return nil, fmt.Errorf("%s: %w", path, err)
if err != nil {
return nil, fmt.Errorf("%s: %w", path, err)
}
return archive.Resources, nil
case "MOD":
// A persistent world never publishes module contents, so the .mod
// contributes no bytes to a manifest — it only says which haks and
// which TLK the manifest covers.
return nil, fmt.Errorf("%s: a module is never emitted; a manifest is haks plus the TLK", path)
} }
} }
switch string(magic) {
case "ERF", "HAK":
index, err := erf.ReadIndex(artifact, size)
if err != nil {
return nil, fmt.Errorf("%s: %w", path, err)
}
return index.Entries, nil
case "MOD":
// A persistent world never publishes module contents, so the .mod
// contributes no bytes to a manifest — it only says which haks and
// which TLK the manifest covers.
return nil, fmt.Errorf("%s: a module is never emitted; a manifest is haks plus the TLK", path)
}
extension := filepath.Ext(filepath.Base(path)) extension := filepath.Ext(filepath.Base(path))
restype, ok := erf.ResourceTypeForExtension(extension) restype, ok := erf.ResourceTypeForExtension(extension)
if !ok { if !ok {
return nil, fmt.Errorf("%s: unknown resource type %q", path, extension) return nil, fmt.Errorf("%s: unknown resource type %q", path, extension)
} }
return []erf.Resource{{ return []erf.IndexEntry{{Name: name, Type: restype, Offset: 0, Size: size}}, nil
Name: name,
Type: restype,
Data: data,
}}, nil
} }
func emitResources(resources []erf.Resource, target sink) ([]Entry, int, int64, error) { // emitResources hashes, compresses and stores one resource at a time, reading
// each payload from the artifact only when its turn comes. Peak memory
// therefore tracks the largest single resource, not the archive: a 2 GB hak
// must emit inside a runner's few spare GB.
func emitResources(artifact io.ReaderAt, index []erf.IndexEntry, target sink) ([]Entry, int, int64, error) {
// A resref appearing twice inside one artifact resolves to the last one, // A resref appearing twice inside one artifact resolves to the last one,
// the way resman lets the last container added win. // the way resman lets the last container added win.
order := make([]Identity, 0, len(resources)) order := make([]Identity, 0, len(index))
latest := make(map[Identity]erf.Resource, len(resources)) latest := make(map[Identity]erf.IndexEntry, len(index))
var tooBig []string var tooBig []string
for _, resource := range resources { for _, entry := range index {
if _, ok := erf.ExtensionForResourceType(resource.Type); !ok { if _, ok := erf.ExtensionForResourceType(entry.Type); !ok {
return nil, 0, 0, fmt.Errorf("resref %s is not resolvable (unknown restype %d)", resource.Name, resource.Type) return nil, 0, 0, fmt.Errorf("resref %s is not resolvable (unknown restype %d)", entry.Name, entry.Type)
} }
if slices.Contains(skippedTypes, resource.Type) { if slices.Contains(skippedTypes, entry.Type) {
continue continue
} }
if len(resource.Data) > fileSizeLimit { if entry.Size > fileSizeLimit {
tooBig = append(tooBig, fmt.Sprintf("%s: %d bytes > %d", resource.Name, len(resource.Data), fileSizeLimit)) tooBig = append(tooBig, fmt.Sprintf("%s: %d bytes > %d", entry.Name, entry.Size, fileSizeLimit))
continue continue
} }
identity := Identity{ResRef: strings.ToLower(resource.Name), ResType: resource.Type} identity := Identity{ResRef: strings.ToLower(entry.Name), ResType: entry.Type}
if _, seen := latest[identity]; !seen { if _, seen := latest[identity]; !seen {
order = append(order, identity) order = append(order, identity)
} }
latest[identity] = resource latest[identity] = entry
} }
if len(tooBig) > 0 { if len(tooBig) > 0 {
sort.Strings(tooBig) sort.Strings(tooBig)
@@ -201,15 +210,18 @@ func emitResources(resources []erf.Resource, target sink) ([]Entry, int, int64,
var blobs int var blobs int
var onDiskBytes int64 var onDiskBytes int64
for _, identity := range order { for _, identity := range order {
resource := latest[identity] payload, err := erf.ReadPayload(artifact, latest[identity])
sum := sha1.Sum(resource.Data) if err != nil {
return nil, 0, 0, err
}
sum := sha1.Sum(payload)
entries = append(entries, Entry{ entries = append(entries, Entry{
SHA1: sum, SHA1: sum,
Size: uint32(len(resource.Data)), Size: uint32(len(payload)),
ResRef: identity.ResRef, ResRef: identity.ResRef,
ResType: identity.ResType, ResType: identity.ResType,
}) })
written, err := target.putBlob(fmt.Sprintf("%x", sum), func() []byte { return compressBlob(resource.Data) }) written, err := target.putBlob(fmt.Sprintf("%x", sum), func() []byte { return compressBlob(payload) })
if err != nil { if err != nil {
return nil, 0, 0, err return nil, 0, 0, err
} }
+126
View File
@@ -0,0 +1,126 @@
package nwsync
import (
"fmt"
"os"
"path/filepath"
"runtime"
"runtime/debug"
"testing"
"time"
"git.westgate.pw/ShadowsOverWestgate/sow-tools/internal/erf"
)
// 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)
}
}
+8 -3
View File
@@ -6,6 +6,7 @@ import (
"crypto/sha256" "crypto/sha256"
"encoding/hex" "encoding/hex"
"fmt" "fmt"
"io"
"os" "os"
"path" "path"
"path/filepath" "path/filepath"
@@ -193,14 +194,18 @@ func resolveIndexKey(artifactKey, outDir string) (string, error) {
// checkArtifactKey fails closed when the key's embedded digest is not the // checkArtifactKey fails closed when the key's embedded digest is not the
// digest of the bytes being emitted. Publishing an index under the wrong key // digest of the bytes being emitted. Publishing an index under the wrong key
// silently pairs a manifest with the wrong artifact. // silently pairs a manifest with the wrong artifact.
func checkArtifactKey(artifactKey string, artifact []byte) error { // artifact is hashed by streaming, so a multi-gigabyte hak is never resident.
func checkArtifactKey(artifactKey string, artifact io.Reader) error {
base := path.Base(artifactKey) base := path.Base(artifactKey)
digest := strings.TrimSuffix(base, path.Ext(base)) digest := strings.TrimSuffix(base, path.Ext(base))
if len(digest) != 64 { if len(digest) != 64 {
return fmt.Errorf("artifact key %q does not name a sha256", artifactKey) return fmt.Errorf("artifact key %q does not name a sha256", artifactKey)
} }
sum := sha256.Sum256(artifact) hash := sha256.New()
if got := hex.EncodeToString(sum[:]); got != digest { if _, err := io.Copy(hash, artifact); err != nil {
return fmt.Errorf("hash artifact: %w", err)
}
if got := hex.EncodeToString(hash.Sum(nil)); got != digest {
return fmt.Errorf("artifact key %q names digest %s but the file hashes to %s", artifactKey, digest, got) return fmt.Errorf("artifact key %q names digest %s but the file hashes to %s", artifactKey, digest, got)
} }
return nil return nil