Files
tensor/io/walk_cycle_pins_test.go
T
petrbalvin af4ee19703
Release / gates (push) Successful in 4m38s
Test / test (push) Successful in 5m16s
Release / release (push) Successful in 35s
feat: initial release
Assisted-by: GLM 5.3 Flash
2026-09-03 10:00:00 +02:00

205 lines
7.1 KiB
Go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT
package io
import (
"encoding/binary"
"fmt"
"math"
"runtime"
"strings"
"testing"
"time"
)
// Regression pins for the HDF5 group walk: deep chains that must stay
// linear, the depth cap, self-links and hard-link diamonds that must
// be refused as cycles.
// groupChain builds a well-formed HDF5 file holding a chain of
// depth groups, each carrying one named attribute and a hard link to the
// next. Every group is legal, the attributes are legal scalar attributes
// and the chain terminates, so the reader has no grounds to refuse it:
// the walk has to stay cheap on its own.
func groupChain(depth int) []byte {
const (
firstGroup = 128
stride = 88
)
n := firstGroup + depth*stride + 16
f := make([]byte, n)
copy(f, hdf5Magic)
f[8] = 0 // superblock version 0
f[13] = 8
f[14] = 8
binary.LittleEndian.PutUint64(f[32:], math.MaxUint64) // free space undefined
binary.LittleEndian.PutUint64(f[40:], uint64(n)) // end of file
binary.LittleEndian.PutUint64(f[48:], math.MaxUint64) // driver undefined
binary.LittleEndian.PutUint64(f[64:], firstGroup) // root object header
for k := range depth {
a := firstGroup + k*stride
nmsg := 2
if k == depth-1 {
nmsg = 1 // the tail group only carries its attribute
}
f[a] = 1 // object header version 1
binary.LittleEndian.PutUint16(f[a+2:], uint16(nmsg))
binary.LittleEndian.PutUint32(f[a+4:], 1) // reference count
binary.LittleEndian.PutUint32(f[a+8:], 88) // message data size
// Attribute message (type 12), size 40, body at a+24.
binary.LittleEndian.PutUint16(f[a+16:], hdf5MsgAttribute)
binary.LittleEndian.PutUint16(f[a+18:], 40)
f[a+24] = 1 // version
binary.LittleEndian.PutUint16(f[a+26:], 8) // name length
binary.LittleEndian.PutUint16(f[a+28:], 8) // datatype length
binary.LittleEndian.PutUint16(f[a+30:], 8) // dataspace length
copy(f[a+32:], fmt.Sprintf("a%07d", k)) // 8-byte attribute name
f[a+40] = 0x10 // fixed-point datatype
binary.LittleEndian.PutUint32(f[a+44:], 1) // one byte
f[a+48] = 1 // scalar dataspace
f[a+56] = byte(k) // one byte of value
if k == depth-1 {
continue
}
// Link message (type 6), size 12, body at a+72.
binary.LittleEndian.PutUint16(f[a+64:], hdf5MsgLink)
binary.LittleEndian.PutUint16(f[a+66:], 12)
f[a+72] = 1 // version
f[a+73] = 0 // flags: hard link, 1-byte name length
f[a+74] = 1 // name length
f[a+75] = 'a'
binary.LittleEndian.PutUint64(f[a+76:], uint64(a+stride))
}
return f
}
// selfLink builds the smallest HDF5 file whose root group links to
// itself: 136 bytes, one object header, one link message.
func selfLink() []byte {
const N = 136
f := make([]byte, N)
copy(f, hdf5Magic)
f[8] = 0 // superblock version 0
f[13] = 8
f[14] = 8
binary.LittleEndian.PutUint64(f[32:], math.MaxUint64)
binary.LittleEndian.PutUint64(f[40:], N)
binary.LittleEndian.PutUint64(f[48:], math.MaxUint64)
binary.LittleEndian.PutUint64(f[64:], 96) // root object header
f[96] = 1 // object header version 1
binary.LittleEndian.PutUint16(f[98:], 1)
binary.LittleEndian.PutUint32(f[100:], 1)
binary.LittleEndian.PutUint32(f[104:], 24)
binary.LittleEndian.PutUint16(f[112:], hdf5MsgLink)
binary.LittleEndian.PutUint16(f[114:], 12)
f[120] = 1 // link message version
f[121] = 0 // hard link, 1-byte name length
f[122] = 1 // name length
f[123] = 'a'
binary.LittleEndian.PutUint64(f[124:], 96) // the link target: itself
return f
}
// TestWalkRefusesHardLinkCycle pins the crash that took the editor down:
// walk had no visited set, so a group linked into itself recursed for
// ever while the path string grew, and the heap grew with it at hundreds
// of megabytes per second until the host ran out of memory. The reader
// must refuse the file with an error.
func TestWalkRefusesHardLinkCycle(t *testing.T) {
guard := time.AfterFunc(20*time.Second, func() { panic("LoadHDF5 on a cyclic file did not return") })
defer guard.Stop()
path := writeHostile(t, "selflink.h5", selfLink())
stop := warnHeap(t, 256<<20)
_, err := LoadHDF5(path)
stop()
if err == nil {
t.Fatal("LoadHDF5 accepted a group that hard-links into itself")
}
if !strings.Contains(err.Error(), "hard-link cycle") {
t.Fatalf("LoadHDF5 = %v, want the cycle refusal", err)
}
}
// TestWalkDeepChainStaysLinear pins the other half of the same crash: a
// well-formed chain of groups used to cost memory quadratic in its
// depth, because every level copied the inherited attribute map and
// built a longer path string. The live memory must grow with the depth,
// not with its square, and the walk carries one shared path buffer and
// one shared attribute map to make that so.
func TestWalkDeepChainStaysLinear(t *testing.T) {
guard := time.AfterFunc(60*time.Second, func() { panic("deep chain walk did not return") })
defer guard.Stop()
measure := func(depth int) (uint64, int) {
data := groupChain(depth)
path := writeHostile(t, fmt.Sprintf("chain%d.h5", depth), data)
runtime.GC()
var before, after runtime.MemStats
runtime.ReadMemStats(&before)
if _, err := LoadHDF5(path); err != nil {
t.Fatalf("depth %d: LoadHDF5 refused a well-formed file: %v", depth, err)
}
runtime.ReadMemStats(&after)
return after.TotalAlloc - before.TotalAlloc, len(data)
}
small, smallFile := measure(150)
big, bigFile := measure(450)
// Linear: three times the depth is about three times the bytes. The
// per-level copies this replaced measured 16.5x for the same step.
if ratio := float64(big) / float64(small); ratio > 6 {
t.Errorf("allocation grows with the square of the depth: %.1fx for 3x depth (%d B for %d B, %d B for %d B)",
ratio, big, bigFile, small, smallFile)
}
}
// TestWalkRefusesUnboundedNesting pins the depth cap: past it the reader
// refuses the file loudly and cheaply instead of walking every level
// first.
func TestWalkRefusesUnboundedNesting(t *testing.T) {
guard := time.AfterFunc(60*time.Second, func() { panic("deep chain walk did not return") })
defer guard.Stop()
deep := writeHostile(t, "deep.h5", groupChain(2*hdf5MaxGroupDepth))
stop := warnHeap(t, 256<<20)
_, err := LoadHDF5(deep)
stop()
if err == nil {
t.Fatalf("LoadHDF5 accepted a chain %d groups deep", 2*hdf5MaxGroupDepth)
}
if !strings.Contains(err.Error(), "nest deeper than") {
t.Fatalf("LoadHDF5 = %v, want the nesting refusal", err)
}
}
// warnHeap fails the test as soon as the live heap passes limit, so a
// regression cannot consume the machine: the watchdog panics, which
// unwinds the offending walk instead of letting it allocate on.
func warnHeap(t *testing.T, limit uint64) func() {
t.Helper()
done := make(chan struct{})
go func() {
ticker := time.NewTicker(20 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-done:
return
case <-ticker.C:
var m runtime.MemStats
runtime.ReadMemStats(&m)
if m.HeapAlloc > limit {
panic("walk heap above its cap")
}
}
}
}()
return func() { close(done) }
}