Skip to content

Commit 4d31265

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 4d31265

8 files changed

Lines changed: 1175 additions & 29 deletions

File tree

‎cmd/osgen/emit/frag_union.go‎

Lines changed: 228 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,173 @@ 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. There is deliberately no Type() method
264+
// or discriminant constants: the wire never identifies the branch.
265+
type {{$t.Name}} struct {
266+
raw json.RawMessage
267+
value any
268+
}
269+
270+
// RawJSON returns the union's JSON bytes. After decoding these are borrowed
271+
// from the response buffer: valid only while the owning response value is
272+
// reachable, must not be mutated, and must be copied if retained beyond it.
273+
func (u *{{$t.Name}}) RawJSON() json.RawMessage { return u.raw }
274+
275+
// SetRaw stages pre-encoded JSON for marshaling.
276+
func (u *{{$t.Name}}) SetRaw(raw json.RawMessage) {
277+
u.raw = raw
278+
u.value = nil
279+
}
280+
{{range $t.Branches}}
281+
// As{{.Name}} decodes the union as {{qualify .GoType}}. The caller selects the
282+
// type it requested; an empty value and nil error mean the union is empty.
283+
func (u *{{$t.Name}}) As{{.Name}}() ({{qualify .GoType}}, error) {
284+
if v, ok := u.value.(*{{qualify .GoType}}); ok {
285+
return *v, nil
286+
}
287+
var v {{qualify .GoType}}
288+
if len(u.raw) == 0 {
289+
return v, nil
290+
}
291+
err := json.Unmarshal(u.raw, &v)
292+
return v, err
293+
}
294+
295+
// New{{$t.Name}}From{{.Name}} returns a {{$t.Name}} populated with v
296+
// on the {{.Name}} branch.
297+
func New{{$t.Name}}From{{.Name}}(v {{qualify .GoType}}) {{$t.Name}} {
298+
return {{$t.Name}}{
299+
value: &v,
300+
}
301+
}
302+
{{end}}
303+
func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
304+
u.raw = data
305+
return nil
306+
}
307+
308+
func (u {{$t.Name}}) MarshalJSON() ([]byte, error) {
309+
if u.value != nil {
310+
return json.Marshal(u.value)
311+
}
312+
if len(u.raw) > 0 {
313+
return u.raw, nil
314+
}
315+
return build.NullJSON, nil
316+
}
317+
{{- else if isTryEach $t.Kind}}
121318
{{- if $t.Comment}}
122319
{{comment $t.Comment}}
123320
{{- else}}
@@ -145,7 +342,9 @@ const (
145342
// Returns {{$t.Name}}UnknownType if the value has not been decoded.
146343
func (u *{{$t.Name}}) Type() {{$t.Name}}Type { return u.typ }
147344
148-
// RawJSON returns the original JSON bytes for escape-hatch decoding.
345+
// RawJSON returns the union's JSON bytes. After decoding these are borrowed
346+
// from the response buffer: valid only while the owning response value is
347+
// reachable, must not be mutated, and must be copied if retained beyond it.
149348
func (u *{{$t.Name}}) RawJSON() json.RawMessage { return u.raw }
150349
151350
// SetRaw stages pre-encoded JSON for marshaling. MarshalJSON emits raw
@@ -160,21 +359,24 @@ func (u *{{$t.Name}}) SetRaw(raw json.RawMessage) {
160359
{{range $t.Branches}}
161360
// {{.Name}} returns the {{qualify .GoType}} branch value.
162361
func (u *{{$t.Name}}) {{.Name}}() {{qualify .GoType}} {
163-
v, _ := u.value.({{qualify .GoType}})
164-
return v
362+
if v, ok := u.value.(*{{qualify .GoType}}); ok {
363+
return *v
364+
}
365+
var zero {{qualify .GoType}}
366+
return zero
165367
}
166368
167369
// New{{$t.Name}}From{{.Name}} returns a {{$t.Name}} populated with v
168370
// on the {{.Name}} branch.
169371
func New{{$t.Name}}From{{.Name}}(v {{qualify .GoType}}) {{$t.Name}} {
170372
return {{$t.Name}}{
171373
typ: {{constName $t.Name .Name}},
172-
value: v,
374+
value: &v,
173375
}
174376
}
175377
{{end}}
176378
func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
177-
u.raw = append(u.raw[:0], data...)
379+
u.raw = data
178380
if len(data) == 0 || bytes.Equal(data, build.NullJSON) {
179381
return nil
180382
}
@@ -189,7 +391,7 @@ func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
189391
var v {{qualify .GoType}}
190392
if err := json.Unmarshal(data, &v); err == nil {
191393
u.typ = {{constName $t.Name .Name}}
192-
u.value = v
394+
u.value = &v
193395
return nil
194396
}
195397
}
@@ -202,7 +404,7 @@ func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
202404
var v {{qualify .GoType}}
203405
if err := json.Unmarshal(data, &v); err == nil {
204406
u.typ = {{constName $t.Name .Name}}
205-
u.value = v
407+
u.value = &v
206408
return nil
207409
}
208410
}
@@ -248,7 +450,9 @@ const (
248450
// Returns {{$t.Name}}UnknownType if the value has not been decoded.
249451
func (u *{{$t.Name}}) Type() {{$t.Name}}Type { return u.typ }
250452
251-
// RawJSON returns the original JSON bytes for escape-hatch decoding.
453+
// RawJSON returns the union's JSON bytes. After decoding these are borrowed
454+
// from the response buffer: valid only while the owning response value is
455+
// reachable, must not be mutated, and must be copied if retained beyond it.
252456
func (u *{{$t.Name}}) RawJSON() json.RawMessage { return u.raw }
253457
254458
// SetRaw stages pre-encoded JSON for marshaling. MarshalJSON emits raw
@@ -263,21 +467,24 @@ func (u *{{$t.Name}}) SetRaw(raw json.RawMessage) {
263467
{{range $t.Branches}}
264468
// {{.Name}} returns the {{qualify .GoType}} branch value.
265469
func (u *{{$t.Name}}) {{.Name}}() {{qualify .GoType}} {
266-
v, _ := u.value.({{qualify .GoType}})
267-
return v
470+
if v, ok := u.value.(*{{qualify .GoType}}); ok {
471+
return *v
472+
}
473+
var zero {{qualify .GoType}}
474+
return zero
268475
}
269476
270477
// New{{$t.Name}}From{{.Name}} returns a {{$t.Name}} populated with v
271478
// on the {{.Name}} branch.
272479
func New{{$t.Name}}From{{.Name}}(v {{qualify .GoType}}) {{$t.Name}} {
273480
return {{$t.Name}}{
274481
typ: {{constName $t.Name .Name}},
275-
value: v,
482+
value: &v,
276483
}
277484
}
278485
{{end}}
279486
func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
280-
u.raw = append(u.raw[:0], data...)
487+
u.raw = data
281488
if len(data) == 0 || bytes.Equal(data, build.NullJSON) {
282489
return nil
283490
}
@@ -290,39 +497,39 @@ func (u *{{$t.Name}}) UnmarshalJSON(data []byte) error {
290497
return err
291498
}
292499
u.typ = {{constName $t.Name .Name}}
293-
u.value = v
500+
u.value = &v
294501
{{- else if eq (tokenStr .TokenClass) "array"}}
295502
case data[0] == '[':
296503
var v {{qualify .GoType}}
297504
if err := json.Unmarshal(data, &v); err != nil {
298505
return err
299506
}
300507
u.typ = {{constName $t.Name .Name}}
301-
u.value = v
508+
u.value = &v
302509
{{- else if eq (tokenStr .TokenClass) "string"}}
303510
case data[0] == '"':
304511
var v {{qualify .GoType}}
305512
if err := json.Unmarshal(data, &v); err != nil {
306513
return err
307514
}
308515
u.typ = {{constName $t.Name .Name}}
309-
u.value = v
516+
u.value = &v
310517
{{- else if eq (tokenStr .TokenClass) "number"}}
311518
case data[0] >= '0' && data[0] <= '9' || data[0] == '-':
312519
var v {{qualify .GoType}}
313520
if err := json.Unmarshal(data, &v); err != nil {
314521
return err
315522
}
316523
u.typ = {{constName $t.Name .Name}}
317-
u.value = v
524+
u.value = &v
318525
{{- else if eq (tokenStr .TokenClass) "bool"}}
319526
case data[0] == 't' || data[0] == 'f':
320527
var v {{qualify .GoType}}
321528
if err := json.Unmarshal(data, &v); err != nil {
322529
return err
323530
}
324531
u.typ = {{constName $t.Name .Name}}
325-
u.value = v
532+
u.value = &v
326533
{{- end}}
327534
{{- end}}
328535
default:

0 commit comments

Comments
 (0)