From d08ccdaaad0ff00622321957aa01cac149741005 Mon Sep 17 00:00:00 2001
From: Zhaozhong Ni <nzz@google.com>
Date: Tue, 28 Aug 2018 09:20:05 -0700
Subject: sentry: avoid double counting map objects in save / restore stats.

PiperOrigin-RevId: 210551929
Change-Id: Idd05935bffc63b39166cc3751139aff61b689faa
---
 pkg/state/decode.go |  6 ++++++
 pkg/state/encode.go |  6 ++++++
 pkg/state/stats.go  | 11 +++++++++++
 3 files changed, 23 insertions(+)

(limited to 'pkg/state')

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 {
-- 
cgit v1.2.3