Add opt-in unmarshal_slab feature: slab-allocated UnmarshalVTSlab with counting prescan

[merlimat]

Motivation

For RPC servers, unmarshalling is where vtprotobuf still spends most of its allocations: every repeated message element, every nested singular message, and every explicit-presence scalar (optional int64, optional string, …) is a separate heap object, so a batch of N puts costs O(N·fields) allocations that all become GC work the moment the request is done. pool helps for the top-level message but does not compose down the tree (children still allocate, and pooling has its own lifecycle burden), and unmarshal_unsafe only addresses string/bytes copies.

The observation behind this PR: for the dominant server pattern — decode a request, handle it, drop it — all those little objects have exactly the same lifetime. They can be carved out of a handful of chunks instead of allocated one by one, without changing the message types, the decode semantics, or the call sites.

This was first validated as a hand-written prototype against real Oxia schemas; two lessons from that prototype shaped the design:

1. Naive chunked slabs are time-neutral at best and regress small messages (chunk zeroing + waste). A counting pre-pass with exact-size reservation is what turns slabs into a win.
2. Even with the pre-pass, tiny payloads pay arena setup without amortization, so small payloads must bypass the arena entirely.

What is generated

Messages opt in per message. For each opted-in root, the generator emits a public entry point; the messages reachable from it (within the same .proto file) get unexported arena-threaded decoders, and the file gets one unexported arena type:

func (m *NestedOrder) UnmarshalVTSlab(dAtA []byte) error {
	if len(dAtA) < protohelpers.SlabUnmarshalThreshold {
		return m.UnmarshalVT(dAtA)
	}
	var a slabArena_slab_nested_proto
	fieldNums := [2]int32{4, 5}
	var counts [2]int
	protohelpers.CountFields(dAtA, fieldNums[:], counts[:])
	if counts[0]+counts[1] < protohelpers.SlabUnmarshalMinCount {
		return m.UnmarshalVT(dAtA)
	}
	a.f_NestedParty.Reserve(1)
	a.f_NestedAddress.Reserve(2)
	// ...
	a.f_NestedLine.Reserve(counts[0])
	a.f_NestedProduct.Reserve(counts[0])
	a.f_NestedNote.Reserve(counts[1])
	return m.unmarshalVTSlab(dAtA, &a)
}

The decode loops are the stock unmarshal loops with only the allocation sites redirected:

• append(m.Puts, &PutRequest{}) becomes append(m.Puts, a.f_PutRequest.Next()), and the element decode threads the arena;
• singular message fields, map values and oneof message payloads come from the same per-type slabs;
• presence pointers (m.Shard = &v) become value copies into per-type scalar slabs (m.Shard = a.f_int64.NextValue(v)).

protohelpers.Slab[T] is a trivial generic bump allocator over []T chunks (reserve-exact, or 8→1024 doubling growth when a reservation is absent or undershoots).

Design

Opt-in, mirroring mempool. A new message option option (vtproto.slab) = true; (extension 64103) plus slab= / slab-exclude= plugin flags, surfaced through GeneratedFile.ShouldSlab. The feature itself is unmarshal_slab, registered the same way unmarshal_unsafe is — a mode of the unmarshal feature, so the emission code is shared rather than forked.

A separate entry point, not a changed UnmarshalVT. UnmarshalVT output is byte-for-byte unchanged; callers that retain messages keep a safe default. Adoption is still zero-effort for the common path because the bundled gRPC and DRPC codecs prefer UnmarshalVTSlab when the message implements it — the schema owner's opt-in is what flips the switch.

Per-file arena, same-file closure. The arena struct and the threaded decoders are generated into the same output file, so the output is self-contained under per-file generation (protoc and buf both). Fields whose types live in other files/packages, well-known types, and groups keep the stock path; unknown-field retention works because the generated code lives in the message's own package.

Counting pre-pass and reservations. One linear scan over the top-level records counts every repeated-message and message-valued-map field of the root in a single pass (CountFields). Reservations are multiplicity-aware: for each counted type, a memoized walk computes how many instances of every type one parent instance can allocate through chains of plain singular message fields, so bounds like Reserve(2 + 3*counts[0]) fall out; the root's own singular chains reserve constants. Oneof variants are excluded from bounds (at most one is set; reserving all would systematically over-allocate) and fall back to growth. Reservations are hints only — undershoot falls back to chunked growth, malformed input just stops the scan early — so there is no correctness coupling between the pre-pass and the decoder.

Two bypasses, both from the prototype's scars. Payloads under SlabUnmarshalThreshold (256 B) skip everything. Payloads whose counted elements sum below SlabUnmarshalMinCount (4) fall back after the scan: a payload can be large yet carry only a couple of elements spread across many types, where per-type chunk minimums cost more than amortization saves (the nested benchmark below was +21% before this bypass, exact stock parity after).

Unchanged semantics. Merge-on-non-fresh, unknown-field retention, required-field tracking, extension ranges, string/bytes copying (no aliasing of the input buffer) are all identical to UnmarshalVT — the emission body is shared, only allocation sites differ. Combining (vtproto.slab) with (vtproto.mempool) on one message is rejected at generation time, since returning chunk-interior pointers to a pool is a footgun.

The trade-off (why this is strictly opt-in)

Slab elements are carved out of shared chunks. A decoded message behaves exactly like an independently allocated one, except that retaining any pointer into the message graph (a sub-message, a repeated element, an optional field pointer) keeps that pointer's whole chunk alive. This is the right trade for messages whose lifetime ends with the RPC that carried them, and the wrong one for messages you dissect and retain pieces of — hence per-message opt-in by the schema owner, and prominent documentation in the README and on the generated method.

Performance

testproto/slab benchmarks, Apple M1 Max, UnmarshalVT vs UnmarshalVTSlab on the same types.

Flat batch (SlabBatch: repeated items + map + presence-heavy fields):

items	time	allocs/op	B/op
8	-22%	151 → 76	+20%
32	-25%	531 → 233	+16%
128	-31%	2035 → 812	+6%

Nested document (NestedOrder: 4–5 message levels, maps, oneofs, optionals at every level):

lines	time	allocs/op	B/op
1	±0 (count bypass → stock)	62 → 62	±0
8	-9%	262 → 160	+15%
32	-15%	978 → 553	+14%
128	-15%	3843 → 2093	+18%

Real-world validation — Oxia's production schemas regenerated with this branch and run through Oxia's serde benchmark harness (count=5 medians):

scenario	time	allocs/op	B/op
WriteRequest, 1 put (128 B)	-2% (bypass)	7 → 7	±0
WriteRequest, 32 puts	-10%	112 → 74	+3%
WriteRequest, 128 puts	-10%	426 → 269	+1%
ReadResponse, 32 gets	-13%	103 → 41	+1%
NotificationBatch, 32	-30%	167 → 47	+12%

Geomean across the engaged Oxia scenarios: -13% decode time, -46% allocs/op. The generated code matches the hand-written prototype's numbers exactly, and the B/op overhead (unused chunk tails) stays within a few percent on realistic shapes.

Testing

• testproto/slab covers repeated/singular/map/oneof/recursive message fields and every presence-pointer scalar kind, across two proto files (two arenas in one package). Tests assert: equivalence with UnmarshalVT on fresh and merged decodes, unknown-field retention through decode + re-encode, no aliasing of the input buffer, no panics on truncation/corruption at every byte offset, and strictly reduced allocations.
• protohelpers unit tests cover slab growth/reservation/undershoot and CountFields (top-level-only counting, group bail-out, truncation, huge phantom lengths).
• All pre-existing generated output (testprotos, WKTs, conformance protos) regenerates byte-identically with the feature enabled — the feature emits nothing unless a message opts in, and UnmarshalVT emission is untouched (so conformance coverage is unaffected).

Limitations / future work

• The slab closure stops at file boundaries; cross-file fields use the stock path. This keeps per-file generation self-contained; extending element-only slabbing across files is possible later.
• Scalar slabs and oneof-reached types are never reserved (growth only); packed repeated scalars and bytes/string contents keep their stock allocation.
• No unmarshal_slab_unsafe combination yet; slab uses the safe string/bytes copies.

[mattlord]

Hi @merlimat !

I see a few issues that I think we should address before merging.

1. I think that we should bound length-delimited records before converting to int in protohelpers/slab.go:91-105. CountFields converts the untrusted uint64 length to int and only checks the resulting index afterward. A counted field followed by MaxUint64-10 makes the index wrap back to zero on 64-bit systems, so a slab-enabled RPC payload can loop forever in the pre-pass. I reproduced this with a two-second test timeout. I think we should reject lengths larger than the remaining buffer before conversion and add this exact overflow case through UnmarshalVTSlab.

2. I think that we should cap reservations derived from possible singular descendants in features/unmarshal/slab.go:283-345. The multiplicity walk eagerly reserves every possible singular descendant even when none appears on the wire. With a legal five-level schema containing ten singular child fields per level, it generated Reserve(100000 * counts[0]); a valid 268-byte payload containing four empty roots allocated roughly 24.5 MB before decoding any children. Since the codecs select this path automatically, that seems like a significant remote memory-amplification/OOM risk. I think we should cap the total reservation using a payload-derived budget, or leave speculative descendants to normal slab growth, and add a sparse branching-schema regression test.

3. I think that we should make unmarshal_slab depend on unmarshal in features/unmarshal/slab.go:305-321. The generator accepts features=unmarshal_slab, but that configuration emits UnmarshalVTSlab without emitting UnmarshalVT, even though both bypass paths call it. The resulting generated package does not compile. I think we should automatically include unmarshal, or reject this feature combination with a clear generator error, and cover the standalone feature selection in a generation test.

Thanks, @merlimat ! ❤️