diff --git a/internal/erf/erf.go b/internal/erf/erf.go index 5e207fc..372878d 100644 --- a/internal/erf/erf.go +++ b/internal/erf/erf.go @@ -311,63 +311,110 @@ func Write(w io.Writer, archive Archive) error { 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: a caller must not mutate +// Data. Callers that only need one resource at a time should use ReadIndex +// instead, which never holds the archive at all. func Read(r io.Reader) (Archive, error) { data, err := io.ReadAll(r) if err != nil { return Archive{}, fmt.Errorf("read erf: %w", err) } - if len(data) < headerSize { - return Archive{}, fmt.Errorf("erf file too small: %d bytes", len(data)) + index, err := ReadIndex(bytes.NewReader(data), int64(len(data))) + if err != nil { + return Archive{}, err } - var hdr header - if err := binary.Read(bytes.NewReader(data[:headerSize]), binary.LittleEndian, &hdr); err != nil { - 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 := make([]Resource, 0, len(index.Entries)) + for _, entry := range index.Entries { resources = append(resources, Resource{ - Name: resref, - Type: key.ResourceType, - Data: payload, - Size: int64(entry.Size), + Name: entry.Name, + Type: entry.Type, + Data: data[entry.Offset : entry.Offset+entry.Size], + Size: entry.Size, }) } return Archive{ - FileType: string(hdr.FileType[:]), - Version: string(hdr.Version[:]), + FileType: index.FileType, + Version: index.Version, Resources: resources, }, nil } diff --git a/internal/nwsync/compressedbuf.go b/internal/nwsync/compressedbuf.go index 43ceca7..af708d3 100644 --- a/internal/nwsync/compressedbuf.go +++ b/internal/nwsync/compressedbuf.go @@ -22,9 +22,13 @@ const ( 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 ( - blobEncoder, _ = zstd.NewWriter(nil) - blobDecoder, _ = zstd.NewReader(nil) + blobEncoder, _ = zstd.NewWriter(nil, zstd.WithEncoderConcurrency(1)) + blobDecoder, _ = zstd.NewReader(nil, zstd.WithDecoderConcurrency(1)) ) // compressBlob wraps data in NWCompressedBuffer framing. diff --git a/internal/nwsync/emit.go b/internal/nwsync/emit.go index df00477..7574e70 100644 --- a/internal/nwsync/emit.go +++ b/internal/nwsync/emit.go @@ -1,10 +1,10 @@ package nwsync import ( - "bytes" "context" "crypto/sha1" "fmt" + "io" "os" "path" "path/filepath" @@ -79,11 +79,18 @@ type EmitOptions struct { // real blobs in the zone and no index, which is unambiguous. Blob names are // content hashes, so re-running skips whatever already landed. 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) } - if err := checkArtifactKey(options.ArtifactKey, artifact); err != nil { + defer artifact.Close() + info, err := artifact.Stat() + if err != nil { + return EmitResult{}, fmt.Errorf("read artifact: %w", err) + } + // A section reader, not the file itself: hashing must not move the file + // offset out from under everything that reads the artifact afterwards. + if err := checkArtifactKey(options.ArtifactKey, io.NewSectionReader(artifact, 0, info.Size())); err != nil { return EmitResult{}, err } name := options.As @@ -98,7 +105,7 @@ func Emit(options EmitOptions) (EmitResult, error) { return EmitResult{}, err } - resources, err := readArtifact(options.ArtifactPath, artifact, name) + index, err := readArtifactIndex(options.ArtifactPath, artifact, info.Size(), name) if err != nil { return EmitResult{}, err } @@ -108,7 +115,7 @@ func Emit(options EmitOptions) (EmitResult, error) { return EmitResult{}, err } - entries, blobs, onDiskBytes, err := emitResources(resources, target) + entries, blobs, onDiskBytes, err := emitResources(artifact, index, target) if err != nil { return EmitResult{}, err } @@ -137,60 +144,64 @@ func openSink(outDir string, injected sink) (sink, error) { return newZoneSink(context.Background(), os.Getenv) } -// readArtifact returns the resources of an ERF/HAK/MOD, or the single resource -// a loose file represents. Upstream's resman does the same dispatch on the -// file's first three bytes. name is the artifact's published name, which for a -// loose file is also its resref. -func readArtifact(path string, data []byte, name string) ([]erf.Resource, error) { - if len(data) >= 3 { - switch string(data[:3]) { - case "ERF", "HAK": - archive, err := erf.Read(bytes.NewReader(data)) - 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) +// readArtifactIndex locates the resources of an ERF/HAK/MOD, or the single +// resource a loose file represents, without reading any payload. Upstream's +// resman does the same dispatch on the file's first three bytes. name is the +// artifact's published name, which for a loose file is also its resref. +func readArtifactIndex(path string, artifact io.ReaderAt, size int64, name string) ([]erf.IndexEntry, error) { + magic := make([]byte, 3) + if size >= 3 { + if _, err := artifact.ReadAt(magic, 0); err != nil { + return nil, fmt.Errorf("%s: %w", path, err) } } + 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)) restype, ok := erf.ResourceTypeForExtension(extension) if !ok { return nil, fmt.Errorf("%s: unknown resource type %q", path, extension) } - return []erf.Resource{{ - Name: name, - Type: restype, - Data: data, - }}, nil + return []erf.IndexEntry{{Name: name, Type: restype, Offset: 0, Size: size}}, 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, // the way resman lets the last container added win. - order := make([]Identity, 0, len(resources)) - latest := make(map[Identity]erf.Resource, len(resources)) + order := make([]Identity, 0, len(index)) + latest := make(map[Identity]erf.IndexEntry, len(index)) var tooBig []string - for _, resource := range resources { - if _, ok := erf.ExtensionForResourceType(resource.Type); !ok { - return nil, 0, 0, fmt.Errorf("resref %s is not resolvable (unknown restype %d)", resource.Name, resource.Type) + for _, entry := range index { + if _, ok := erf.ExtensionForResourceType(entry.Type); !ok { + 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 } - if len(resource.Data) > fileSizeLimit { - tooBig = append(tooBig, fmt.Sprintf("%s: %d bytes > %d", resource.Name, len(resource.Data), fileSizeLimit)) + if entry.Size > fileSizeLimit { + tooBig = append(tooBig, fmt.Sprintf("%s: %d bytes > %d", entry.Name, entry.Size, fileSizeLimit)) 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 { order = append(order, identity) } - latest[identity] = resource + latest[identity] = entry } if len(tooBig) > 0 { sort.Strings(tooBig) @@ -201,15 +212,18 @@ func emitResources(resources []erf.Resource, target sink) ([]Entry, int, int64, var blobs int var onDiskBytes int64 for _, identity := range order { - resource := latest[identity] - sum := sha1.Sum(resource.Data) + payload, err := erf.ReadPayload(artifact, latest[identity]) + if err != nil { + return nil, 0, 0, err + } + sum := sha1.Sum(payload) entries = append(entries, Entry{ SHA1: sum, - Size: uint32(len(resource.Data)), + Size: uint32(len(payload)), ResRef: identity.ResRef, 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 { return nil, 0, 0, err } diff --git a/internal/nwsync/memory_test.go b/internal/nwsync/memory_test.go new file mode 100644 index 0000000..1a94d08 --- /dev/null +++ b/internal/nwsync/memory_test.go @@ -0,0 +1,150 @@ +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) + } +} diff --git a/internal/nwsync/sink.go b/internal/nwsync/sink.go index f6ce812..f552737 100644 --- a/internal/nwsync/sink.go +++ b/internal/nwsync/sink.go @@ -6,6 +6,7 @@ import ( "crypto/sha256" "encoding/hex" "fmt" + "io" "os" "path" "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 // digest of the bytes being emitted. Publishing an index under the wrong key // 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) digest := strings.TrimSuffix(base, path.Ext(base)) if len(digest) != 64 { return fmt.Errorf("artifact key %q does not name a sha256", artifactKey) } - sum := sha256.Sum256(artifact) - if got := hex.EncodeToString(sum[:]); got != digest { + hash := sha256.New() + 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 nil