sentry: avoid double counting map objects in save / restore stats.

PiperOrigin-RevId: 210551929
Change-Id: Idd05935bffc63b39166cc3751139aff61b689faa

3 files changed, 23 insertions, 0 deletions

diff --git a/pkg/state/decode.go b/pkg/state/decode.go
index b3fbc793c..3ef59610b 100644
--- a/pkg/state/decode.go
+++ b/pkg/state/decode.go
@@ -365,6 +365,12 @@ func (ds *decodeState) decodeObject(os *objectState, obj reflect.Value, object *
 			// (For non-interfaces this is a no-op).
 			dyntyp := reflect.TypeOf(obj.Interface())
 			if dyntyp.Kind() == reflect.Map {
+				// Remove the map object count here to avoid
+				// double counting, as this object will be
+				// counted again when it gets processed later.
+				// We do not add a reference count as the
+				// reference is artificial.
+				ds.stats.Remove(obj)
 				obj.Set(ds.register(id, dyntyp).obj)
 			} else if dyntyp.Kind() == reflect.Ptr {
 				ds.push(true /* dereference */, "", nil)
diff --git a/pkg/state/encode.go b/pkg/state/encode.go
index 59914bef4..fd052db12 100644
--- a/pkg/state/encode.go
+++ b/pkg/state/encode.go
@@ -335,6 +335,12 @@ func (es *encodeState) encodeObject(obj reflect.Value, mapAsValue bool, format s
 			object = &pb.Object{Value: &pb.Object_MapValue{es.encodeMap(obj)}}
 		} else {
 			// Encode a reference to the map.
+			//
+			// Remove the map object count here to avoid double
+			// counting, as this object will be counted again when
+			// it gets processed later. We do not add a reference
+			// count as the reference is artificial.
+			es.stats.Remove(obj)
 			object = &pb.Object{Value: &pb.Object_RefValue{es.register(obj)}}
 		}
 	default:
diff --git a/pkg/state/stats.go b/pkg/state/stats.go
index c4135a889..ddcc49f78 100644
--- a/pkg/state/stats.go
+++ b/pkg/state/stats.go
@@ -68,6 +68,17 @@ func (s *Stats) Add(obj reflect.Value) {
 	entry.count++
 }
 
+// Remove removes a sample count. It should only be called after a previous
+// Add().
+func (s *Stats) Remove(obj reflect.Value) {
+	if s == nil {
+		return
+	}
+	typ := obj.Type()
+	entry := s.byType[typ]
+	entry.count--
+}
+
 // Start starts a sample.
 func (s *Stats) Start(obj reflect.Value) {
 	if s == nil {