Skip to content

Commit 5941fd8

Browse files
sean-sean-
committed
feat(osgen): single-pass union decode for error unions, lazy As<T>() for aggregations
Two classes of try-each discriminated union are replaced with single-pass strategies. On mget this cuts decode allocations ~2.7x (1000 docs: ~43k -> ~16k allocs/op) and time ~1.7x (4.2ms -> 2.5ms); the remaining cost is the GetResult decode itself plus interface boxing, not the union machinery. - Case A (merged): object unions with one permissive "primary" branch plus discriminated branch(es) -- mget, msearch, indices-open. The primary is embedded and the common case decodes in a single json.Unmarshal; each discriminated branch is detected by the presence of its distinguishing key and decoded only when matched. This drops the build.HasJSONKeys map probe (which was ~61% of the old allocations) and the per-item raw copy. - Case B (lazy As<T>()): aggregation/suggest result unions carry no wire discriminator (avg/sum/min/max all serialize as {"value":N}, and bucket types collide), so they cannot be auto-selected. UnmarshalJSON only retains the raw bytes; generated As<ConcreteType>() accessors decode on demand into the type the caller requested. - Unions fitting neither (e.g. reindex bodies, plugin-defined task status) keep the existing try-each decoder; the classifier logs once per union name when it declines to convert a wrapper-shaped union. - All union UnmarshalJSON now aliases the owned response buffer (u.raw = data) rather than copying it; RawJSON() documents the borrowed-buffer contract (valid while the response is reachable, copy to retain). Signed-off-by: Sean Chittenden <sean.chittenden@crowdstrike.com>
1 parent a9d02f7 commit 5941fd8

7 files changed

Lines changed: 1063 additions & 24 deletions

File tree

‎cmd/osgen/emit/frag_union.go‎

Lines changed: 244 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -24,17 +24,37 @@ type UnionFragment struct {
2424
Registry *ir.TypeRegistry
2525
}
2626

27-
// Imports returns the imports the union-types fragment needs.
27+
// Imports returns the imports the union-types fragment needs. fmt is only used
28+
// by the try-each/first-byte variants (fmt.Errorf); bytes only by the variants
29+
// that null-check with bytes.Equal (everything except the lazy-accessor one).
30+
// build.HasJSONKeys is only emitted for try-each unions.
2831
func (f *UnionFragment) Imports() []Import {
2932
if len(f.Types) == 0 {
3033
return nil
3134
}
3235
imps := []Import{
33-
{Path: "bytes"},
3436
{Path: "encoding/json"},
35-
{Path: "fmt"},
3637
{Path: LocalModule + "/internal/build"},
3738
}
39+
var needBytes, needFmt bool
40+
for _, t := range f.Types {
41+
switch {
42+
case t.Merge != nil:
43+
needBytes = true
44+
case t.LazyAccessors:
45+
// json + build only
46+
case t.Kind == ir.TypeLazyUnion: // try-each
47+
needBytes, needFmt = true, true
48+
default: // first-byte switch
49+
needBytes, needFmt = true, true
50+
}
51+
}
52+
if needBytes {
53+
imps = append(imps, Import{Path: "bytes"})
54+
}
55+
if needFmt {
56+
imps = append(imps, Import{Path: "fmt"})
57+
}
3858
if f.Op != nil && f.Op.IsPlugin && f.Registry != nil && f.hasCrossPkgBranch() {
3959
imps = append(imps, Import{Path: f.Registry.CoreImport})
4060
}
@@ -76,6 +96,7 @@ func (f *UnionFragment) Body() (string, error) {
7696
"isTryEach": func(k ir.TypeKind) bool { return k == ir.TypeLazyUnion },
7797
"qualify": qualify,
7898
"quotedKeys": quotedKeys,
99+
"embedField": embedFieldName,
79100
}).Parse(unionFragTmplText))
80101

81102
if err := tmpl.Execute(&sb, f.Types); err != nil {
@@ -105,6 +126,16 @@ func unionConstNameIR(unionName, branchName string) string {
105126
return unionName + branchName + "Type"
106127
}
107128

129+
// embedFieldName returns the selector used to reference an embedded type: the
130+
// substring after the last package qualifier dot. "opensearchapi.GetResult"
131+
// -> "GetResult"; "GetResult" -> "GetResult".
132+
func embedFieldName(goType string) string {
133+
if i := strings.LastIndex(goType, "."); i >= 0 {
134+
return goType[i+1:]
135+
}
136+
return goType
137+
}
138+
108139
// quotedKeys renders a slice of field names as a comma-separated list of
109140
// Go double-quoted string literals, for splicing into a build.HasJSONKeys
110141
// call in the generated try-each discriminator.
@@ -117,7 +148,189 @@ func quotedKeys(keys []string) string {
117148
}
118149

119150
const unionFragTmplText = `{{- range $t := .}}
120-
{{- if isTryEach $t.Kind}}
151+
{{- if $t.Merge}}
152+
{{- if $t.Comment}}
153+
{{comment $t.Comment}}
154+
{{- else}}
155+
// {{$t.Name}} is a discriminated union type (single-pass merge decode).
156+
{{- end}}
157+
// Use Type() to determine which branch was decoded, then call
158+
// the corresponding accessor.
159+
type {{$t.Name}} struct {
160+
typ {{$t.Name}}Type
161+
raw json.RawMessage
162+
value any
163+
}
164+
165+
// {{$t.Name}}Type discriminates the branches of {{$t.Name}}.
166+
type {{$t.Name}}Type int
167+
168+
const (
169+
{{$t.Name}}UnknownType {{$t.Name}}Type = iota
170+
{{- range $t.Branches}}
171+
{{constName $t.Name .Name}}
172+
{{- end}}
173+
)
174+
175+
// Type returns which union branch was populated during decoding.
176+
// Returns {{$t.Name}}UnknownType if the value has not been decoded.
177+
func (u *{{$t.Name}}) Type() {{$t.Name}}Type { return u.typ }
178+
179+
// RawJSON returns the union's JSON bytes. After decoding these are borrowed
180+
// from the response buffer: valid only while the owning response value is
181+
// reachable, must not be mutated, and must be copied if retained beyond it.
182+
func (u *{{$t.Name}}) RawJSON() json.RawMessage { return u.raw }
183+
184+
// SetRaw stages pre-encoded JSON for marshaling. MarshalJSON emits raw
185+
// verbatim when no typed branch is set. Use the New{{$t.Name}}From*
186+
// constructors to populate a typed branch instead; SetRaw is the typed
187+
// escape hatch for callers that already have wire-format bytes.
188+
func (u *{{$t.Name}}) SetRaw(raw json.RawMessage) {
189+
u.raw = raw
190+
u.value = nil
191+
u.typ = {{$t.Name}}UnknownType
192+
}
193+
{{range $t.Branches}}
194+
// {{.Name}} returns the {{qualify .GoType}} branch value.
195+
func (u *{{$t.Name}}) {{.Name}}() {{qualify .GoType}} {
196+
if v, ok := u.value.(*{{qualify .GoType}}); ok {
197+
return *v
198+
}
199+
var zero {{qualify .GoType}}
200+
return zero
201+
}
202+
203+
// New{{$t.Name}}From{{.Name}} returns a {{$t.Name}} populated with v
204+
// on the {{.Name}} branch.
205+
func New{{$t.Name}}From{{.Name}}(v {{qualify .GoType}}) {{$t.Name}} {
206+
return {{$t.Name}}{
207+
typ: {{constName $t.Name .Name}},
208+
value: &v,
209+
}
210+
}
211+
{{end}}
212+
func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
213+
u.raw = data
214+
if len(data) == 0 || bytes.Equal(data, build.NullJSON) {
215+
return nil
216+
}
217+
// Single decode: embed the permissive (primary) branch and probe for the
218+
// discriminating keys of the other branches in one pass. encoding/json
219+
// populates the embedded primary directly; the probes only test presence.
220+
type merged struct {
221+
{{qualify $t.Merge.PrimaryGoType}}
222+
{{- range $t.Merge.Probes}}
223+
{{.GoName}} json.RawMessage ` + "`json:\"{{.JSONKey}}\"`" + `
224+
{{- end}}
225+
}
226+
var m merged
227+
if err := json.Unmarshal(data, &m); err != nil {
228+
return err
229+
}
230+
{{- range $t.Merge.Branches}}
231+
if {{range $i, $p := .PresentProbes}}{{if $i}} && {{end}}len(m.{{$p}}) > 0{{end}} {
232+
var v {{qualify .GoType}}
233+
if err := json.Unmarshal(data, &v); err != nil {
234+
return err
235+
}
236+
u.typ = {{.Const}}
237+
u.value = &v
238+
return nil
239+
}
240+
{{- end}}
241+
u.typ = {{$t.Merge.PrimaryConst}}
242+
u.value = &m.{{embedField (qualify $t.Merge.PrimaryGoType)}}
243+
return nil
244+
}
245+
246+
func (u {{$t.Name}}) MarshalJSON() ([]byte, error) {
247+
if u.value != nil {
248+
return json.Marshal(u.value)
249+
}
250+
if len(u.raw) > 0 {
251+
return u.raw, nil
252+
}
253+
return build.NullJSON, nil
254+
}
255+
{{- else if $t.LazyAccessors}}
256+
{{- if $t.Comment}}
257+
{{comment $t.Comment}}
258+
{{- else}}
259+
// {{$t.Name}} is a discriminated union with no wire discriminator.
260+
{{- end}}
261+
// Its branches are indistinguishable from the response bytes alone (the type
262+
// is determined by the request), so the raw JSON is retained and decoded on
263+
// demand by the As<Branch>() accessors.
264+
type {{$t.Name}} struct {
265+
typ {{$t.Name}}Type
266+
raw json.RawMessage
267+
value any
268+
}
269+
270+
// {{$t.Name}}Type discriminates the branches of {{$t.Name}}.
271+
type {{$t.Name}}Type int
272+
273+
const (
274+
{{$t.Name}}UnknownType {{$t.Name}}Type = iota
275+
{{- range $t.Branches}}
276+
{{constName $t.Name .Name}}
277+
{{- end}}
278+
)
279+
280+
// Type returns the branch set by a constructor, or {{$t.Name}}UnknownType for a
281+
// decoded value (the wire does not identify the branch; use an As* accessor).
282+
func (u *{{$t.Name}}) Type() {{$t.Name}}Type { return u.typ }
283+
284+
// RawJSON returns the union's JSON bytes. After decoding these are borrowed
285+
// from the response buffer: valid only while the owning response value is
286+
// reachable, must not be mutated, and must be copied if retained beyond it.
287+
func (u *{{$t.Name}}) RawJSON() json.RawMessage { return u.raw }
288+
289+
// SetRaw stages pre-encoded JSON for marshaling.
290+
func (u *{{$t.Name}}) SetRaw(raw json.RawMessage) {
291+
u.raw = raw
292+
u.value = nil
293+
u.typ = {{$t.Name}}UnknownType
294+
}
295+
{{range $t.Branches}}
296+
// As{{.Name}} decodes the union as {{qualify .GoType}}. The caller selects the
297+
// type it requested; an empty value and nil error mean the union is empty.
298+
func (u *{{$t.Name}}) As{{.Name}}() ({{qualify .GoType}}, error) {
299+
if v, ok := u.value.(*{{qualify .GoType}}); ok {
300+
return *v, nil
301+
}
302+
var v {{qualify .GoType}}
303+
if len(u.raw) == 0 {
304+
return v, nil
305+
}
306+
err := json.Unmarshal(u.raw, &v)
307+
return v, err
308+
}
309+
310+
// New{{$t.Name}}From{{.Name}} returns a {{$t.Name}} populated with v
311+
// on the {{.Name}} branch.
312+
func New{{$t.Name}}From{{.Name}}(v {{qualify .GoType}}) {{$t.Name}} {
313+
return {{$t.Name}}{
314+
typ: {{constName $t.Name .Name}},
315+
value: &v,
316+
}
317+
}
318+
{{end}}
319+
func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
320+
u.raw = data
321+
return nil
322+
}
323+
324+
func (u {{$t.Name}}) MarshalJSON() ([]byte, error) {
325+
if u.value != nil {
326+
return json.Marshal(u.value)
327+
}
328+
if len(u.raw) > 0 {
329+
return u.raw, nil
330+
}
331+
return build.NullJSON, nil
332+
}
333+
{{- else if isTryEach $t.Kind}}
121334
{{- if $t.Comment}}
122335
{{comment $t.Comment}}
123336
{{- else}}
@@ -145,7 +358,9 @@ const (
145358
// Returns {{$t.Name}}UnknownType if the value has not been decoded.
146359
func (u *{{$t.Name}}) Type() {{$t.Name}}Type { return u.typ }
147360
148-
// RawJSON returns the original JSON bytes for escape-hatch decoding.
361+
// RawJSON returns the union's JSON bytes. After decoding these are borrowed
362+
// from the response buffer: valid only while the owning response value is
363+
// reachable, must not be mutated, and must be copied if retained beyond it.
149364
func (u *{{$t.Name}}) RawJSON() json.RawMessage { return u.raw }
150365
151366
// SetRaw stages pre-encoded JSON for marshaling. MarshalJSON emits raw
@@ -160,21 +375,24 @@ func (u *{{$t.Name}}) SetRaw(raw json.RawMessage) {
160375
{{range $t.Branches}}
161376
// {{.Name}} returns the {{qualify .GoType}} branch value.
162377
func (u *{{$t.Name}}) {{.Name}}() {{qualify .GoType}} {
163-
v, _ := u.value.({{qualify .GoType}})
164-
return v
378+
if v, ok := u.value.(*{{qualify .GoType}}); ok {
379+
return *v
380+
}
381+
var zero {{qualify .GoType}}
382+
return zero
165383
}
166384
167385
// New{{$t.Name}}From{{.Name}} returns a {{$t.Name}} populated with v
168386
// on the {{.Name}} branch.
169387
func New{{$t.Name}}From{{.Name}}(v {{qualify .GoType}}) {{$t.Name}} {
170388
return {{$t.Name}}{
171389
typ: {{constName $t.Name .Name}},
172-
value: v,
390+
value: &v,
173391
}
174392
}
175393
{{end}}
176394
func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
177-
u.raw = append(u.raw[:0], data...)
395+
u.raw = data
178396
if len(data) == 0 || bytes.Equal(data, build.NullJSON) {
179397
return nil
180398
}
@@ -189,7 +407,7 @@ func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
189407
var v {{qualify .GoType}}
190408
if err := json.Unmarshal(data, &v); err == nil {
191409
u.typ = {{constName $t.Name .Name}}
192-
u.value = v
410+
u.value = &v
193411
return nil
194412
}
195413
}
@@ -202,7 +420,7 @@ func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
202420
var v {{qualify .GoType}}
203421
if err := json.Unmarshal(data, &v); err == nil {
204422
u.typ = {{constName $t.Name .Name}}
205-
u.value = v
423+
u.value = &v
206424
return nil
207425
}
208426
}
@@ -248,7 +466,9 @@ const (
248466
// Returns {{$t.Name}}UnknownType if the value has not been decoded.
249467
func (u *{{$t.Name}}) Type() {{$t.Name}}Type { return u.typ }
250468
251-
// RawJSON returns the original JSON bytes for escape-hatch decoding.
469+
// RawJSON returns the union's JSON bytes. After decoding these are borrowed
470+
// from the response buffer: valid only while the owning response value is
471+
// reachable, must not be mutated, and must be copied if retained beyond it.
252472
func (u *{{$t.Name}}) RawJSON() json.RawMessage { return u.raw }
253473
254474
// SetRaw stages pre-encoded JSON for marshaling. MarshalJSON emits raw
@@ -263,21 +483,24 @@ func (u *{{$t.Name}}) SetRaw(raw json.RawMessage) {
263483
{{range $t.Branches}}
264484
// {{.Name}} returns the {{qualify .GoType}} branch value.
265485
func (u *{{$t.Name}}) {{.Name}}() {{qualify .GoType}} {
266-
v, _ := u.value.({{qualify .GoType}})
267-
return v
486+
if v, ok := u.value.(*{{qualify .GoType}}); ok {
487+
return *v
488+
}
489+
var zero {{qualify .GoType}}
490+
return zero
268491
}
269492
270493
// New{{$t.Name}}From{{.Name}} returns a {{$t.Name}} populated with v
271494
// on the {{.Name}} branch.
272495
func New{{$t.Name}}From{{.Name}}(v {{qualify .GoType}}) {{$t.Name}} {
273496
return {{$t.Name}}{
274497
typ: {{constName $t.Name .Name}},
275-
value: v,
498+
value: &v,
276499
}
277500
}
278501
{{end}}
279502
func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
280-
u.raw = append(u.raw[:0], data...)
503+
u.raw = data
281504
if len(data) == 0 || bytes.Equal(data, build.NullJSON) {
282505
return nil
283506
}
@@ -290,39 +513,39 @@ func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
290513
return err
291514
}
292515
u.typ = {{constName $t.Name .Name}}
293-
u.value = v
516+
u.value = &v
294517
{{- else if eq (tokenStr .TokenClass) "array"}}
295518
case data[0] == '[':
296519
var v {{qualify .GoType}}
297520
if err := json.Unmarshal(data, &v); err != nil {
298521
return err
299522
}
300523
u.typ = {{constName $t.Name .Name}}
301-
u.value = v
524+
u.value = &v
302525
{{- else if eq (tokenStr .TokenClass) "string"}}
303526
case data[0] == '"':
304527
var v {{qualify .GoType}}
305528
if err := json.Unmarshal(data, &v); err != nil {
306529
return err
307530
}
308531
u.typ = {{constName $t.Name .Name}}
309-
u.value = v
532+
u.value = &v
310533
{{- else if eq (tokenStr .TokenClass) "number"}}
311534
case data[0] >= '0' && data[0] <= '9' || data[0] == '-':
312535
var v {{qualify .GoType}}
313536
if err := json.Unmarshal(data, &v); err != nil {
314537
return err
315538
}
316539
u.typ = {{constName $t.Name .Name}}
317-
u.value = v
540+
u.value = &v
318541
{{- else if eq (tokenStr .TokenClass) "bool"}}
319542
case data[0] == 't' || data[0] == 'f':
320543
var v {{qualify .GoType}}
321544
if err := json.Unmarshal(data, &v); err != nil {
322545
return err
323546
}
324547
u.typ = {{constName $t.Name .Name}}
325-
u.value = v
548+
u.value = &v
326549
{{- end}}
327550
{{- end}}
328551
default:

0 commit comments

Comments
 (0)