Files
sow-tools/internal/nwsync/sink.go
T
archvillainetteandClaude Opus 5 59b233e3db fix(nwsync): stream emit so peak memory tracks the largest resource (#76)
Emit read the whole artifact into memory and erf.Read then allocated a
second full copy of every payload, so a hak cost about 5x its size in
RAM. The 7 GB runner was OOM-killed on any hak over ~1.4 GB, which
blocks the NWSync backfill and every release that rebuilds a large hak.

Emit now opens the artifact, hashes it by streaming for the key check,
parses only the header and resource table via erf.ReadIndex, and reads,
hashes, compresses and stores one payload at a time. erf.Read keeps its
old shape but returns payloads as subslices instead of fresh copies,
which removes the second copy for the other callers too.

The zstd encoder pool also held one window-sized history per CPU — about
200 MB of live heap on a 24-core runner. EncodeAll is single-threaded
per call, so concurrency 1 gives byte-identical blobs for far less
memory.

Peak heap is now flat at ~22 MB for both an 8 MB and a 64 MB hak, and a
regression test asserts it does not scale with artifact size.

Closes #76

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-07-31 13:06:29 +02:00

213 lines
7.3 KiB
Go

package nwsync
import (
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"io"
"os"
"path"
"path/filepath"
"strings"
"git.westgate.pw/ShadowsOverWestgate/sow-tools/internal/depot"
)
// sink is where an emit or assemble run puts what it produces. The zone is the
// production sink; a local directory exists only as the conformance path, so
// upstream's output and ours can be diffed on a developer machine.
type sink interface {
// putBlob stores one NWCompressedBuffer blob under its sha1 name and
// returns the bytes stored, or 0 if the blob was already there. Blob names
// are content hashes, so an existing name is existing content — which is
// why body is a thunk: compression is the expensive part of emit and a
// blob that is already stored must not pay for it.
putBlob(sha1Hex string, body func() []byte) (int64, error)
// putIndex stores a NSYM manifest and its sidecar under key, which is
// either an artifact-derived object key or a local path.
putIndex(key string, manifest, sidecar []byte) error
// getIndex reads back a NSYM manifest and its sidecar.
getIndex(key string) (manifest, sidecar []byte, err error)
// describe names the sink for messages.
describe(key string) string
}
// dirSink writes a local NWSync repository tree.
type dirSink struct{ root string }
func (s dirSink) putBlob(sha1Hex string, body func() []byte) (int64, error) {
blob := blobPath(s.root, sha1Hex)
if _, err := os.Stat(blob); err == nil {
return 0, nil
}
if err := os.MkdirAll(filepath.Dir(blob), 0o755); err != nil {
return 0, fmt.Errorf("create blob directory: %w", err)
}
data := body()
if err := os.WriteFile(blob, data, 0o644); err != nil {
return 0, fmt.Errorf("write blob: %w", err)
}
return int64(len(data)), nil
}
func (s dirSink) putIndex(key string, manifest, sidecar []byte) error {
target := filepath.Join(s.root, filepath.FromSlash(key))
if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
return fmt.Errorf("create manifest directory: %w", err)
}
if err := os.WriteFile(target, manifest, 0o644); err != nil {
return fmt.Errorf("write manifest: %w", err)
}
if err := os.WriteFile(target+".json", sidecar, 0o644); err != nil {
return fmt.Errorf("write sidecar: %w", err)
}
return nil
}
func (s dirSink) getIndex(key string) ([]byte, []byte, error) {
target := filepath.Join(s.root, filepath.FromSlash(key))
manifest, err := os.ReadFile(target)
if err != nil {
return nil, nil, fmt.Errorf("read index: %w", err)
}
sidecar, err := os.ReadFile(target + ".json")
if err != nil {
return nil, nil, fmt.Errorf("read sidecar: %w", err)
}
return manifest, sidecar, nil
}
func (s dirSink) describe(key string) string {
return filepath.Join(s.root, filepath.FromSlash(key))
}
// zoneSink uploads straight to the NWSync storage zone. Nothing bulky is ever
// written to the runner's disk: the working set is one resource at a time.
type zoneSink struct {
store depot.KeyStore
ctx context.Context
zone string
}
func (s zoneSink) putBlob(sha1Hex string, body func() []byte) (int64, error) {
key := path.Join("data", "sha1", sha1Hex[0:2], sha1Hex[2:4], sha1Hex)
// A throttled probe must never be read as "missing, re-upload" or as
// "present, skip", so only a confirmed Present skips the upload.
state, _, err := s.store.ProbeKey(s.ctx, key)
if err != nil {
return 0, fmt.Errorf("probe blob %s: %w", sha1Hex, err)
}
if state == depot.Present {
return 0, nil
}
data := body()
if err := s.put(key, data); err != nil {
return 0, fmt.Errorf("upload blob %s: %w", sha1Hex, err)
}
return int64(len(data)), nil
}
func (s zoneSink) putIndex(key string, manifest, sidecar []byte) error {
// The manifest lands last: its presence is the publication marker, so it
// must never appear before the blobs it names.
if err := s.put(key+".json", sidecar); err != nil {
return fmt.Errorf("upload sidecar %s: %w", key, err)
}
if err := s.put(key, manifest); err != nil {
return fmt.Errorf("upload index %s: %w", key, err)
}
return nil
}
func (s zoneSink) getIndex(key string) ([]byte, []byte, error) {
manifest, err := s.store.GetKey(s.ctx, key)
if err != nil {
return nil, nil, fmt.Errorf("read index: %w", err)
}
sidecar, err := s.store.GetKey(s.ctx, key+".json")
if err != nil {
return nil, nil, fmt.Errorf("read sidecar: %w", err)
}
return manifest, sidecar, nil
}
func (s zoneSink) describe(key string) string { return s.zone + "/" + key }
func (s zoneSink) put(key string, body []byte) error {
sum := sha256.Sum256(body)
return s.store.PutReader(s.ctx, key, bytes.NewReader(body), int64(len(body)), hex.EncodeToString(sum[:]))
}
// newZoneSink builds the upload sink from the environment. NWSync data lives
// in its own zone, separate from the asset depot, so it has its own zone and
// credential; only the host is shared, and Crucible has no default host.
func newZoneSink(ctx context.Context, getenv func(string) string) (sink, error) {
cfg := depot.LoadConfig(getenv)
cfg.StorageZone = getenv("NWSYNC_STORAGE_ZONE")
cfg.WriteKey = getenv("NWSYNC_STORAGE_PASSWORD")
cfg.ReadKey = cfg.WriteKey
if cfg.StorageZone == "" {
return nil, fmt.Errorf("NWSYNC_STORAGE_ZONE is unset (or pass --out DIR to write locally)")
}
if cfg.WriteKey == "" {
return nil, fmt.Errorf("NWSYNC_STORAGE_PASSWORD is unset (or pass --out DIR to write locally)")
}
if cfg.StorageHost == "" {
return nil, fmt.Errorf("BUNNY_STORAGE_HOST is unset")
}
store, err := depot.NewKeyStore(cfg)
if err != nil {
return nil, err
}
return zoneSink{store: store, ctx: ctx, zone: cfg.StorageZone}, nil
}
// indexKey is where an artifact's NSYM lives: beside the artifact itself, with
// the final extension replaced. emit and assemble must agree on this one rule,
// so it lives here and nowhere else.
//
// artifacts/haks/sha256/30/46/3046….hak -> artifacts/haks/sha256/30/46/3046….nsym
func indexKey(artifactKey string) (string, error) {
extension := path.Ext(artifactKey)
if extension == "" {
return "", fmt.Errorf("artifact key %q has no extension", artifactKey)
}
return strings.TrimSuffix(artifactKey, extension) + ".nsym", nil
}
// resolveIndexKey is where emit writes an artifact's index and where assemble
// reads it from. On the zone that is beside the artifact; locally the indexes
// sit flat beside the data tree, so upstream's output and ours diff directly.
func resolveIndexKey(artifactKey, outDir string) (string, error) {
key, err := indexKey(artifactKey)
if err != nil {
return "", err
}
if outDir != "" {
return path.Base(key), nil
}
return key, nil
}
// 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.
// 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)
}
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
}