From ce52c31877d1a6c041be34321daeafd8f35b1537 Mon Sep 17 00:00:00 2001
From: Damien Mathieu <42@dmathieu.com>
Date: Wed, 26 Jul 2023 11:04:43 +0200
Subject: [PATCH] Upgrade modeljson ints to int64 if their proto uses int64
 (#122)

* introduce a nullable.Int64

* upgrade modeljson int to int64 if their proto uses int64

* remove transforms from int64 to int

* handle int64 in generator

* populate int64
---
 .../internal/modeldecoder/generator/code.go   |  2 +-
 .../modeldecoder/generator/jsonschema.go      |  2 +-
 .../modeldecoder/generator/nullable.go        |  1 +
 .../modeldecoder/generator/preprocessor.go    |  2 +-
 .../modeldecodertest/populator.go             |  3 ++
 .../modeldecoder/nullable/nullable.go         | 34 +++++++++++++++++
 .../modeldecoder/nullable/nullable_test.go    | 38 +++++++++++++++++++
 .../internal/modeldecoder/v2/model.go         |  2 +-
 model/internal/modeljson/marshal_fastjson.go  |  2 +-
 model/internal/modeljson/metricset.go         |  2 +-
 model/internal/modeljson/session.go           |  2 +-
 model/modelpb/apmevent.pb.json.go             |  2 +-
 model/modelpb/span.pb.json.go                 |  4 +-
 model/modelpb/transaction.pb.json.go          |  2 +-
 14 files changed, 87 insertions(+), 11 deletions(-)

diff --git a/input/elasticapm/internal/modeldecoder/generator/code.go b/input/elasticapm/internal/modeldecoder/generator/code.go
index ab4003c0..da27684b 100644
--- a/input/elasticapm/internal/modeldecoder/generator/code.go
+++ b/input/elasticapm/internal/modeldecoder/generator/code.go
@@ -350,7 +350,7 @@ if !val.IsSet() {
 		switch f.Type().String() {
 		case nullableTypeString:
 			validation = generateNullableStringValidation
-		case nullableTypeInt:
+		case nullableTypeInt, nullableTypeInt64:
 			validation = generateNullableIntValidation
 		case nullableTypeFloat64:
 			// right now we can reuse the validation rules for int
diff --git a/input/elasticapm/internal/modeldecoder/generator/jsonschema.go b/input/elasticapm/internal/modeldecoder/generator/jsonschema.go
index cd34225b..6aed2475 100644
--- a/input/elasticapm/internal/modeldecoder/generator/jsonschema.go
+++ b/input/elasticapm/internal/modeldecoder/generator/jsonschema.go
@@ -112,7 +112,7 @@ func (g *JSONSchemaGenerator) generate(st structType, key string, prop *property
 			err = generateJSONPropertyJSONNumber(&info, prop, &childProp)
 		case nullableTypeHTTPHeader:
 			err = generateJSONPropertyHTTPHeader(&info, prop, &childProp)
-		case nullableTypeInt, nullableTypeTimeMicrosUnix:
+		case nullableTypeInt, nullableTypeInt64, nullableTypeTimeMicrosUnix:
 			err = generateJSONPropertyInteger(&info, prop, &childProp)
 		case nullableTypeInterface:
 			err = generateJSONPropertyInterface(&info, prop, &childProp)
diff --git a/input/elasticapm/internal/modeldecoder/generator/nullable.go b/input/elasticapm/internal/modeldecoder/generator/nullable.go
index cba5cbe1..b5177ea7 100644
--- a/input/elasticapm/internal/modeldecoder/generator/nullable.go
+++ b/input/elasticapm/internal/modeldecoder/generator/nullable.go
@@ -28,6 +28,7 @@ var (
 	nullableTypeFloat64        = fmt.Sprintf("%s.Float64", typPath)
 	nullableTypeHTTPHeader     = fmt.Sprintf("%s.HTTPHeader", typPath)
 	nullableTypeInt            = fmt.Sprintf("%s.Int", typPath)
+	nullableTypeInt64          = fmt.Sprintf("%s.Int64", typPath)
 	nullableTypeInterface      = fmt.Sprintf("%s.Interface", typPath)
 	nullableTypeString         = fmt.Sprintf("%s.String", typPath)
 	nullableTypeTimeMicrosUnix = fmt.Sprintf("%s.TimeMicrosUnix", typPath)
diff --git a/input/elasticapm/internal/modeldecoder/generator/preprocessor.go b/input/elasticapm/internal/modeldecoder/generator/preprocessor.go
index 968991eb..d4be6ceb 100644
--- a/input/elasticapm/internal/modeldecoder/generator/preprocessor.go
+++ b/input/elasticapm/internal/modeldecoder/generator/preprocessor.go
@@ -70,7 +70,7 @@ func validateJSONTag(flatTag, nestedTag string) bool {
 
 func getFlatFieldType(f structField) (string, error) {
 	switch typ := f.Type().String(); typ {
-	case nullableTypeInt:
+	case nullableTypeInt, nullableTypeInt64:
 		return "json.Number", nil
 	case nullableTypeString:
 		return "string", nil
diff --git a/input/elasticapm/internal/modeldecoder/modeldecodertest/populator.go b/input/elasticapm/internal/modeldecoder/modeldecodertest/populator.go
index d8092f75..187a6b42 100644
--- a/input/elasticapm/internal/modeldecoder/modeldecodertest/populator.go
+++ b/input/elasticapm/internal/modeldecoder/modeldecodertest/populator.go
@@ -182,6 +182,9 @@ func SetStructValues(in interface{}, values *Values, opts ...SetStructValuesOpti
 			case nullable.Int:
 				v.Set(values.Int)
 				fieldVal = reflect.ValueOf(v)
+			case nullable.Int64:
+				v.Set(int64(values.Int))
+				fieldVal = reflect.ValueOf(v)
 			case nullable.Interface:
 				if strings.Contains(key, "port") {
 					v.Set(values.Int)
diff --git a/input/elasticapm/internal/modeldecoder/nullable/nullable.go b/input/elasticapm/internal/modeldecoder/nullable/nullable.go
index d7c8c428..97295567 100644
--- a/input/elasticapm/internal/modeldecoder/nullable/nullable.go
+++ b/input/elasticapm/internal/modeldecoder/nullable/nullable.go
@@ -54,6 +54,15 @@ func init() {
 			(*((*Int)(ptr))).isSet = true
 		}
 	})
+	jsoniter.RegisterTypeDecoderFunc("nullable.Int64", func(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+		switch iter.WhatIsNext() {
+		case jsoniter.NilValue:
+			iter.ReadNil()
+		default:
+			(*((*Int64)(ptr))).Val = iter.ReadInt64()
+			(*((*Int64)(ptr))).isSet = true
+		}
+	})
 	jsoniter.RegisterTypeDecoderFunc("nullable.Float64", func(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
 		switch iter.WhatIsNext() {
 		case jsoniter.NilValue:
@@ -193,6 +202,31 @@ func (v *Int) Reset() {
 	v.isSet = false
 }
 
+// Int64 stores an int64 value and the
+// information if the value has been set
+type Int64 struct {
+	Val   int64
+	isSet bool
+}
+
+// Set sets the value
+func (v *Int64) Set(val int64) {
+	v.Val = val
+	v.isSet = true
+}
+
+// IsSet is true when decode was called
+func (v *Int64) IsSet() bool {
+	return v.isSet
+}
+
+// Reset sets the Int64 to it's initial state
+// where it is not set and has no value
+func (v *Int64) Reset() {
+	v.Val = 0
+	v.isSet = false
+}
+
 // Float64 stores a float64 value and the
 // information if the value has been set
 type Float64 struct {
diff --git a/input/elasticapm/internal/modeldecoder/nullable/nullable_test.go b/input/elasticapm/internal/modeldecoder/nullable/nullable_test.go
index 4d803fa2..71bed051 100644
--- a/input/elasticapm/internal/modeldecoder/nullable/nullable_test.go
+++ b/input/elasticapm/internal/modeldecoder/nullable/nullable_test.go
@@ -31,6 +31,7 @@ import (
 type testType struct {
 	S   String         `json:"s"`
 	I   Int            `json:"i"`
+	I64 Int64          `json:"i64"`
 	F   Float64        `json:"f"`
 	B   Bool           `json:"b"`
 	V   Interface      `json:"v"`
@@ -114,6 +115,43 @@ func TestInt(t *testing.T) {
 	}
 }
 
+func TestInt64(t *testing.T) {
+	for _, tc := range []struct {
+		name  string
+		input string
+
+		val         int64
+		isSet, fail bool
+	}{
+		{name: "values", input: `{"i64":44}`, val: 44, isSet: true},
+		{name: "empty", input: `{"i64":0}`, isSet: true},
+		{name: "null", input: `{"i64":null}`, isSet: false},
+		{name: "missing", input: `{}`},
+		{name: "invalid", input: `{"i64":"1.0.1"}`, fail: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			dec := json.NewDecoder(strings.NewReader(tc.input))
+			var testStruct testType
+			err := dec.Decode(&testStruct)
+			if tc.fail {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tc.isSet, testStruct.I64.IsSet())
+				assert.Equal(t, tc.val, testStruct.I64.Val)
+			}
+
+			testStruct.I64.Reset()
+			assert.False(t, testStruct.I64.IsSet())
+			assert.Empty(t, testStruct.I64.Val)
+
+			testStruct.I64.Set(55)
+			assert.True(t, testStruct.I64.IsSet())
+			assert.Equal(t, int64(55), testStruct.I64.Val)
+		})
+	}
+}
+
 func TestFloat64(t *testing.T) {
 	for _, tc := range []struct {
 		name  string
diff --git a/input/elasticapm/internal/modeldecoder/v2/model.go b/input/elasticapm/internal/modeldecoder/v2/model.go
index 234263f6..4e4b3286 100644
--- a/input/elasticapm/internal/modeldecoder/v2/model.go
+++ b/input/elasticapm/internal/modeldecoder/v2/model.go
@@ -1117,7 +1117,7 @@ type transactionUserExperience struct {
 
 type longtaskMetrics struct {
 	// Count is the total number of of longtasks.
-	Count nullable.Int `json:"count" validate:"required,min=0"`
+	Count nullable.Int64 `json:"count" validate:"required,min=0"`
 	// Max longtask duration
 	Max nullable.Float64 `json:"max" validate:"required,min=0"`
 	// Sum of longtask durations
diff --git a/model/internal/modeljson/marshal_fastjson.go b/model/internal/modeljson/marshal_fastjson.go
index a6ac41fb..eed7b50a 100644
--- a/model/internal/modeljson/marshal_fastjson.go
+++ b/model/internal/modeljson/marshal_fastjson.go
@@ -2574,7 +2574,7 @@ func (v *Session) MarshalFastJSON(w *fastjson.Writer) error {
 	w.String(v.ID)
 	if v.Sequence != 0 {
 		w.RawString(",\"sequence\":")
-		w.Int64(int64(v.Sequence))
+		w.Int64(v.Sequence)
 	}
 	w.RawByte('}')
 	return nil
diff --git a/model/internal/modeljson/metricset.go b/model/internal/modeljson/metricset.go
index 7442fa59..c799cf90 100644
--- a/model/internal/modeljson/metricset.go
+++ b/model/internal/modeljson/metricset.go
@@ -100,7 +100,7 @@ func (ms *MetricsetSample) MarshalFastJSON(w *fastjson.Writer) error {
 }
 
 type AggregatedDuration struct {
-	Count int
+	Count int64
 	Sum   time.Duration
 }
 
diff --git a/model/internal/modeljson/session.go b/model/internal/modeljson/session.go
index a2b19a88..b45689fe 100644
--- a/model/internal/modeljson/session.go
+++ b/model/internal/modeljson/session.go
@@ -19,5 +19,5 @@ package modeljson
 
 type Session struct {
 	ID       string `json:"id"`
-	Sequence int    `json:"sequence,omitempty"`
+	Sequence int64  `json:"sequence,omitempty"`
 }
diff --git a/model/modelpb/apmevent.pb.json.go b/model/modelpb/apmevent.pb.json.go
index 0fe268f6..5a924429 100644
--- a/model/modelpb/apmevent.pb.json.go
+++ b/model/modelpb/apmevent.pb.json.go
@@ -251,7 +251,7 @@ func (e *APMEvent) MarshalFastJSON(w *fastjson.Writer) error {
 	if e.Session != nil {
 		doc.Session = &modeljson.Session{
 			ID:       e.Session.Id,
-			Sequence: int(e.Session.Sequence),
+			Sequence: e.Session.Sequence,
 		}
 	}
 
diff --git a/model/modelpb/span.pb.json.go b/model/modelpb/span.pb.json.go
index 18576bba..a7ba8ccc 100644
--- a/model/modelpb/span.pb.json.go
+++ b/model/modelpb/span.pb.json.go
@@ -61,7 +61,7 @@ func (e *Span) toModelJSON(out *modeljson.Span) {
 	}
 	if e.SelfTime != nil {
 		out.SelfTime = modeljson.AggregatedDuration{
-			Count: int(e.SelfTime.Count),
+			Count: e.SelfTime.Count,
 			Sum:   e.SelfTime.Sum.AsDuration(),
 		}
 	}
@@ -87,7 +87,7 @@ func (e *Span) toModelJSON(out *modeljson.Span) {
 		}
 		if e.DestinationService.ResponseTime != nil {
 			out.Destination.Service.ResponseTime = modeljson.AggregatedDuration{
-				Count: int(e.DestinationService.ResponseTime.Count),
+				Count: e.DestinationService.ResponseTime.Count,
 				Sum:   e.DestinationService.ResponseTime.Sum.AsDuration(),
 			}
 		}
diff --git a/model/modelpb/transaction.pb.json.go b/model/modelpb/transaction.pb.json.go
index d55ae0d5..f727987f 100644
--- a/model/modelpb/transaction.pb.json.go
+++ b/model/modelpb/transaction.pb.json.go
@@ -71,7 +71,7 @@ func (e *Transaction) toModelJSON(out *modeljson.Transaction, metricset bool) {
 
 				if dss.Duration != nil {
 					dssJson.Duration = modeljson.AggregatedDuration{
-						Count: int(dss.Duration.Count),
+						Count: dss.Duration.Count,
 						Sum:   dss.Duration.Sum.AsDuration(),
 					}
 				}