Memory audit: cut retention and allocation across the stack

README.md

@@ -139,7 +139,7 @@ and maintains it as a project of his own.
 | **XSD 1.1** | 99.98% instance (26,217 of 26,222); **99.97%** schema-validity (15,350 of 15,354); opt-in via `Version11` |
 | **RELAX NG** | 100.00% of James Clark's spectest (965 of 965 assertions); XML and compact syntax, from Go and from `go-xml validate -rng` |
 | **DTD** | content models, attribute defaults, enumerations, `ID`/`IDREF`; external subset, parameter entities across both subsets, conditional sections — via `dtd.Load` with a caller-supplied resolver, nothing fetched by default |
-| **Tests** | 2,854 `func Test` declarations, clean under `-race` (a few subtests skip without the corpora below) |
+| **Tests** | 2,950 `func Test` declarations, clean under `-race` (a few subtests skip without the corpora below) |
 | **Production schemas** | UBL 2.1, UN/CEFACT CII, Factur-X/ZUGFeRD, Peppol BIS 3.0 — 88 schemas load, instances validate clean |
 | **API** | 1.2; the exported surface is stable and additive over 1.1, and a breaking change means 2.0 with a new module path |
 
@@ -669,7 +669,7 @@ exact signal.
 
 `FileResolver`'s document cache exists because `fn:doc` is defined to return
 the *same* node for the same URI within one execution, so `doc('x') is doc('x')`
-requires it. It is mutex-guarded, and now capped at 256 documents. It had been
+requires it. It is mutex-guarded, and now capped at 256 documents and 32 MB of source. It had been
 unbounded on the reasoning that the directories the caller opened with
 `-allow-dir` bound it in practice — true for the size of the *set*, but the
 process holds every document it has ever parsed for its whole lifetime, and a
@@ -1388,7 +1388,7 @@ back, is in [docs/testing.md](docs/testing.md).
 
 | method | what it catches | what it misses |
 |---|---|---|
-| **Unit tests** (2,854 `func Test` declarations) | places where a plausible implementation is quietly wrong | anything nobody thought to write a test for |
+| **Unit tests** (2,950 `func Test` declarations) | places where a plausible implementation is quietly wrong | anything nobody thought to write a test for |
 | **Spec inventories** | features absent entirely | features present but behaving wrongly |
 | **Saxon differential** | subtle behavioural divergence on real stylesheets | constructs the corpora do not use |
 | **W3C QT3 suite** | systematic conformance across 22,054 XPath and 30,517 XQuery cases | XSLT (it is an XPath suite) |

c14n/bytes_alloc_test.go

@@ -0,0 +1,47 @@
+package c14n
+
+import (
+	"runtime"
+	"testing"
+)
+
+// Bytes and BytesNodeSet grew a bytes.Buffer by doubling from empty, 3.7
+// times the canonical form in all, then kept the writer's chunks and joined
+// them, 2.1 times, and Equal paid it for its first operand. They now size
+// the result from the node set and write into it once, about 1.1 times.
+func TestBytesAllocation(t *testing.T) {
+	// 3 MB: doubling cost 3.1x the output here and 3.7x at 10 MB.
+	root := cachedDoc(t, 3<<20).FirstChild()
+	out, err := Bytes(root, Options{Algorithm: Inclusive10})
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, tc := range []struct {
+		name string
+		f    func()
+	}{
+		{"Bytes", func() { _, _ = Bytes(root, Options{Algorithm: Inclusive10}) }},
+		{"BytesNodeSet", func() { _, _ = BytesNodeSet(Subtree(root), Options{Algorithm: Inclusive10}) }},
+		{"Equal", func() { _, _ = Equal(root, root, Options{Algorithm: Inclusive10}) }},
+	} {
+		got := heapAllocated(tc.f)
+		t.Logf("%s: %d bytes allocated for %d bytes of output (%.2fx)", tc.name, got, len(out), float64(got)/float64(len(out)))
+		if got > uint64(len(out))*7/5 {
+			t.Errorf("%s allocated %d bytes for %d bytes of output, want under 1.4x", tc.name, got, len(out))
+		}
+	}
+}
+
+// heapAllocated is the heap f allocates, the least of three runs so that a
+// collection landing inside one does not count against it.
+func heapAllocated(f func()) uint64 {
+	best := ^uint64(0)
+	for range 3 {
+		var a, b runtime.MemStats
+		runtime.ReadMemStats(&a)
+		f()
+		runtime.ReadMemStats(&b)
+		best = min(best, b.TotalAlloc-a.TotalAlloc)
+	}
+	return best
+}

c14n/c14n.go

@@ -192,9 +192,95 @@ func FormatPrefixList(prefixes []string) string {
 // descendants are, with the namespace and xml:* axis treatment the
 // chosen algorithm specifies for a subset apex.
 func Bytes(n *xdm.Node, opts Options) ([]byte, error) {
-	var b bytes.Buffer
-	err := Write(&b, n, opts)
-	return b.Bytes(), err
+	if n == nil || (n.Kind() != xdm.KindDocument && n.Kind() != xdm.KindElement) {
+		return nil, ErrUnsupportedNode
+	}
+	return BytesNodeSet(Subtree(n), opts)
+}
+
+// chunkWriter keeps a copy of everything written to it.
+type chunkWriter struct{ chunks [][]byte }
+
+func (c *chunkWriter) Write(p []byte) (int, error) {
+	c.chunks = append(c.chunks, bytes.Clone(p))
+	return len(p), nil
+}
+
+// bytes joins the chunks. No output is nil, as from an empty bytes.Buffer.
+func (c *chunkWriter) bytes() []byte {
+	switch len(c.chunks) {
+	case 0:
+		return nil
+	case 1:
+		return c.chunks[0]
+	}
+	return bytes.Join(c.chunks, nil)
+}
+
+// sizedWriter collects everything written to it in one buffer, sized once
+// from the tree under the node set's root at the first write (after the options were validated, so
+// a refused call allocates nothing). The estimate runs 5% to 15% over the
+// output, so the result costs about 1.1 times its size where a buffer grown
+// by doubling cost 2.1 times and joined chunks twice; a short estimate (an
+// escaped value is longer than its text) costs one more copy.
+type sizedWriter struct {
+	ns   NodeSet
+	opts Options
+	buf  []byte
+}
+
+func (s *sizedWriter) Write(p []byte) (int, error) {
+	if s.buf == nil {
+		s.buf = make([]byte, 0, max(estimateSize(s.ns, s.opts), len(p)))
+	}
+	if need := len(s.buf) + len(p); need > cap(s.buf) {
+		// A short estimate: grow by a quarter, not by append's doubling.
+		s.buf = append(make([]byte, 0, need+need/4), s.buf...)
+	}
+	s.buf = append(s.buf, p...)
+	return len(p), nil
+}
+
+// estimateSize adds up the markup and values of ns's members. It reads the
+// tree once, a sixth of the time canonicalizing it takes. It counts every
+// declaration an element has and no escaping, and adds an eighth for the
+// escapes, so it is an estimate. It does not ask ns which nodes are members
+// (a NodeSet is asked once per node), so it is for the sets this package
+// made: a Subtree is counted exactly, an ExcludeSubtree with what it excludes.
+func estimateSize(ns NodeSet, opts Options) int {
+	root := ns.Root()
+	if root == nil {
+		return 0
+	}
+	comments := opts.Algorithm.WithComments()
+	n := 0
+	visit := func(d *xdm.Node) {
+		switch d.Kind() {
+		case xdm.KindElement:
+			q := d.Name()
+			n += 2*(len(q.Local)+len(q.Prefix)) + 6
+			for a := range d.Attrs() {
+				q := a.Name()
+				n += len(q.Local) + len(q.Prefix) + len(a.Value()) + 5
+			}
+			for p, u := range d.DeclaredNamespaces() {
+				n += len(p) + len(u) + 9
+			}
+		case xdm.KindComment:
+			if comments {
+				n += len(d.Value()) + 8
+			}
+		case xdm.KindPI:
+			n += len(d.Name().Local) + len(d.Value()) + 8
+		case xdm.KindText:
+			n += len(d.Value())
+		}
+	}
+	visit(root)
+	for d := range root.Descendants() {
+		visit(d)
+	}
+	return n + n/8
 }
 
 // Digest canonicalizes n and writes the result into h, returning h's sum.
@@ -277,9 +363,19 @@ func WriteNodeSet(w io.Writer, ns NodeSet, opts Options) error {
 
 // BytesNodeSet mirrors Bytes for a subset.
 func BytesNodeSet(ns NodeSet, opts Options) ([]byte, error) {
-	var b bytes.Buffer
-	err := WriteNodeSet(&b, ns, opts)
-	return b.Bytes(), err
+	switch ns.(type) {
+	case subtree, excludeSet:
+		c := sizedWriter{ns: ns, opts: opts}
+		err := WriteNodeSet(&c, ns, opts)
+		return c.buf, err
+	}
+	// Any other set is the caller's, and Contains is called once per node:
+	// the output cannot be sized without calling it again, so it arrives in
+	// chunks of about outBufSize, kept and joined once, about twice the
+	// canonical form in all.
+	var c chunkWriter
+	err := WriteNodeSet(&c, ns, opts)
+	return c.bytes(), err
 }
 
 // DigestNodeSet mirrors Digest for a subset.

c14n/canon.go

@@ -6,6 +6,7 @@ import (
 	"regexp"
 	"slices"
 	"strings"
+	"sync"
 
 	"github.com/knroy/go-xml/v2/xdm"
 )
@@ -96,7 +97,10 @@ type canon struct {
 const noNode = "\x00"
 
 func canonicalize(w io.Writer, ns NodeSet, opts Options) error {
-	_, err := run(w, ns, opts)
+	c, err := run(w, ns, opts)
+	if c != nil {
+		c.w.release()
+	}
 	return err
 }
 
@@ -108,7 +112,6 @@ func run(w io.Writer, ns NodeSet, opts Options) (*canon, error) {
 	}
 	_, all := ns.(subtree)
 	c := &canon{
-		w:        newOutBuf(w),
 		set:      ns,
 		all:      all,
 		excl:     opts.Algorithm.Exclusive(),
@@ -139,6 +142,7 @@ func run(w io.Writer, ns NodeSet, opts Options) (*canon, error) {
 			c.scope.m[prefix] = uri
 		}
 	}
+	c.w = newOutBuf(w)
 	var err error
 	switch root.Kind() {
 	case xdm.KindDocument:
@@ -150,6 +154,7 @@ func run(w io.Writer, ns NodeSet, opts Options) (*canon, error) {
 			c.leaf(root)
 		}
 	default:
+		c.w.release()
 		return nil, ErrUnsupportedNode
 	}
 	if err == nil {
@@ -170,12 +175,32 @@ func run(w io.Writer, ns NodeSet, opts Options) (*canon, error) {
 // into nothing.
 type outBuf struct {
 	buf []byte
+	pb  *[]byte // the pool entry buf came from
 	w   io.Writer
 	err error
 }
 
+// outBufPool keeps the buffers between calls: a fixed 80 KiB allocated per
+// call was 81 KB of garbage for each Digest of a tiny document.
+var outBufPool = sync.Pool{New: func() any {
+	b := make([]byte, 0, outBufSize+outBufSize/4)
+	return &b
+}}
+
 func newOutBuf(w io.Writer) *outBuf {
-	return &outBuf{buf: make([]byte, 0, outBufSize+outBufSize/4), w: w}
+	pb := outBufPool.Get().(*[]byte)
+	return &outBuf{buf: (*pb)[:0], pb: pb, w: w}
+}
+
+// release returns the buffer to the pool unless output grew it far past its
+// size, so that one large text node does not stay allocated. o is not used
+// afterwards.
+func (o *outBuf) release() {
+	if cap(o.buf) <= 2*outBufSize {
+		*o.pb = o.buf[:0]
+		outBufPool.Put(o.pb)
+	}
+	o.buf, o.pb = nil, nil
 }
 
 // flush hands buf to w. A short write without an error is io.ErrShortWrite,

c14n/digest_alloc_test.go

@@ -0,0 +1,27 @@
+package c14n
+
+import (
+	"crypto/sha256"
+	"testing"
+)
+
+// Every call allocated an 80 KiB output buffer, 81 KB of garbage for each
+// Digest of a tiny document. The buffers now come from a pool. The bound is half the unpooled
+// figure because the race detector makes a pool drop a quarter of its puts.
+func TestDigestSmallDocumentAllocation(t *testing.T) {
+	root := parse(t, `<a><b>x</b></a>`).FirstChild()
+	h := sha256.New()
+	const calls = 101
+	got := heapAllocated(func() {
+		for range calls {
+			h.Reset()
+			if _, err := Digest(h, root, Options{Algorithm: Inclusive10}); err != nil {
+				t.Fatal(err)
+			}
+		}
+	})
+	t.Logf("%d Digest calls allocated %d bytes", calls, got)
+	if got > 4<<20 {
+		t.Errorf("%d Digest calls allocated %d bytes, want under 4 MiB (8.3 MB with a fresh buffer each)", calls, got)
+	}
+}

cmd/go-xml/gc_test.go

@@ -23,3 +23,23 @@ func TestLowerGCRespectsGOGC(t *testing.T) {
 		t.Errorf("with GOGC=100: GC percent %d, want the environment's 100", got)
 	}
 }
+
+// Past 128 MB of live heap the command drops from GOGC=200 to 100; below it
+// the setting is left alone.
+func TestRelaxGCOnBigLiveHeap(t *testing.T) {
+	orig := debug.SetGCPercent(cliGCPercent)
+	defer debug.SetGCPercent(orig)
+
+	if relaxGC(cliGCBigHeap) {
+		t.Error("relaxed at exactly the threshold, want only past it")
+	}
+	if got := debug.SetGCPercent(cliGCPercent); got != cliGCPercent {
+		t.Errorf("a small heap changed GC percent to %d", got)
+	}
+	if !relaxGC(cliGCBigHeap + 1) {
+		t.Error("did not relax past the threshold")
+	}
+	if got := debug.SetGCPercent(cliGCPercent); got != 100 {
+		t.Errorf("GC percent after a big heap = %d, want 100", got)
+	}
+}

cmd/go-xml/main.go

@@ -11,8 +11,11 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"runtime"
 	"runtime/debug"
+	"runtime/metrics"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/knroy/go-xml/v2/xdm"
@@ -51,12 +54,51 @@ func main() {
 // never in the library, whose callers own their process's GC policy.
 const cliGCPercent = 200
 
+// cliGCBigHeap is the live heap, in bytes, past which the command drops to
+// GOGC=100. At 200 the heap may reach three times what is live, and past this
+// size that is what costs: peak RSS on DocBook table-cals.049 was 1,160 MB at
+// 200 and 789 MB at 100, and on a 100 MB identity transform 1,840-1,956 MB
+// against 1,360-1,400, for the same wall time and 5-10% more CPU. Below it a
+// small heap is cheap to hold and the extra CPU is what 200 saves. The size of
+// the input is no guide -- DocBook's is small -- so it is the live heap that
+// is read.
+const cliGCBigHeap = 128 << 20
+
 // lowerGC applies cliGCPercent unless GOGC is set, so a user who tunes the
-// collector keeps their setting.
+// collector keeps their setting, and watches the live heap to relax it.
 func lowerGC() {
 	if os.Getenv("GOGC") == "" {
 		debug.SetGCPercent(cliGCPercent)
+		watchLiveHeap.Do(armGCSentinel)
+	}
+}
+
+var watchLiveHeap sync.Once
+
+type gcSentinel struct{ _ [16]byte }
+
+// armGCSentinel runs relaxGC after each collection until it has acted. A
+// finalizer on an unreachable object runs once the cycle that found it has
+// marked, so the live-heap metric it reads is that cycle's.
+func armGCSentinel() {
+	runtime.SetFinalizer(new(gcSentinel), func(*gcSentinel) {
+		s := []metrics.Sample{{Name: "/gc/heap/live:bytes"}}
+		metrics.Read(s)
+		if s[0].Value.Kind() == metrics.KindUint64 && relaxGC(s[0].Value.Uint64()) {
+			return
+		}
+		armGCSentinel()
+	})
+}
+
+// relaxGC drops to GOGC=100 when live bytes are past cliGCBigHeap, and
+// reports whether it did.
+func relaxGC(live uint64) bool {
+	if live <= cliGCBigHeap {
+		return false
 	}
+	debug.SetGCPercent(100)
+	return true
 }
 
 // paramFlag collects repeated -p name=value flags.