3 files changed, 30 insertions, 12 deletions
diff --git a/pkg/state/decode.go b/pkg/state/decode.go
index 33ec926c7..b3fbc793c 100644
--- a/pkg/state/decode.go
+++ b/pkg/state/decode.go
@@ -78,12 +78,11 @@ func (os *objectState) checkComplete(stats *Stats) {
 	if os.blockedBy > 0 {
 		return
 	}
+	stats.Start(os.obj)
 
 	// Fire all callbacks.
 	for _, fn := range os.callbacks {
-		stats.Start(os.obj)
 		fn()
-		stats.Done()
 	}
 	os.callbacks = nil
 
@@ -93,6 +92,7 @@ func (os *objectState) checkComplete(stats *Stats) {
 		other.checkComplete(stats)
 	}
 	os.blocking = nil
+	stats.Done()
 }
 
 // waitFor queues a dependency on the given object.
@@ -329,6 +329,7 @@ func (ds *decodeState) decodeInterface(os *objectState, obj reflect.Value, i *pb
 // decodeObject decodes a object value.
 func (ds *decodeState) decodeObject(os *objectState, obj reflect.Value, object *pb.Object, format string, param interface{}) {
 	ds.push(false, format, param)
+	ds.stats.Add(obj)
 	ds.stats.Start(obj)
 
 	switch x := object.GetValue().(type) {
@@ -466,12 +467,14 @@ func (ds *decodeState) Deserialize(obj reflect.Value) {
 	// See above, we never process objects while we have no outstanding
 	// interests (other than the very first object).
 	for id := uint64(1); ds.outstanding > 0; id++ {
+		os := ds.lookup(id)
+		ds.stats.Start(os.obj)
+
 		o, err := ds.readObject()
 		if err != nil {
 			panic(err)
 		}
 
-		os := ds.lookup(id)
 		if os != nil {
 			// Decode the object.
 			ds.from = &os.path
@@ -483,6 +486,8 @@ func (ds *decodeState) Deserialize(obj reflect.Value) {
 			// registered.
 			ds.deferred[id] = o
 		}
+
+		ds.stats.Done()
 	}
 
 	// Check the zero-length header at the end.
diff --git a/pkg/state/encode.go b/pkg/state/encode.go
index 1cec14f24..59914bef4 100644
--- a/pkg/state/encode.go
+++ b/pkg/state/encode.go
@@ -241,6 +241,7 @@ func (es *encodeState) encodeInterface(obj reflect.Value) *pb.Interface {
 // If mapAsValue is true, then a map will be encoded directly.
 func (es *encodeState) encodeObject(obj reflect.Value, mapAsValue bool, format string, param interface{}) (object *pb.Object) {
 	es.push(false, format, param)
+	es.stats.Add(obj)
 	es.stats.Start(obj)
 
 	switch obj.Kind() {
@@ -354,10 +355,13 @@ func (es *encodeState) Serialize(obj reflect.Value) {
 	// Pop off the list until we're done.
 	for es.pending.Len() > 0 {
 		e := es.pending.Front()
-		es.pending.Remove(e)
 
 		// Extract the queued object.
 		qo := e.Value.(queuedObject)
+		es.stats.Start(qo.obj)
+
+		es.pending.Remove(e)
+
 		es.from = &qo.path
 		o := es.encodeObject(qo.obj, true, "", nil)
 
@@ -368,6 +372,7 @@ func (es *encodeState) Serialize(obj reflect.Value) {
 
 		// Mark as done.
 		es.done.PushBack(e)
+		es.stats.Done()
 	}
 
 	// Write a zero-length terminal at the end; this is a sanity check
diff --git a/pkg/state/stats.go b/pkg/state/stats.go
index 1ebd8ebb4..c4135a889 100644
--- a/pkg/state/stats.go
+++ b/pkg/state/stats.go
@@ -44,20 +44,28 @@ type Stats struct {
 	last time.Time
 }
 
-// sample adds the given number of samples to the given object.
-func (s *Stats) sample(typ reflect.Type, count uint) {
+// sample adds the samples to the given object.
+func (s *Stats) sample(typ reflect.Type) {
+	now := time.Now()
+	s.byType[typ].total += now.Sub(s.last)
+	s.last = now
+}
+
+// Add adds a sample count.
+func (s *Stats) Add(obj reflect.Value) {
+	if s == nil {
+		return
+	}
 	if s.byType == nil {
 		s.byType = make(map[reflect.Type]*statEntry)
 	}
+	typ := obj.Type()
 	entry, ok := s.byType[typ]
 	if !ok {
 		entry = new(statEntry)
 		s.byType[typ] = entry
 	}
-	now := time.Now()
-	entry.count += count
-	entry.total += now.Sub(s.last)
-	s.last = now
+	entry.count++
 }
 
 // Start starts a sample.
@@ -67,7 +75,7 @@ func (s *Stats) Start(obj reflect.Value) {
 	}
 	if len(s.stack) > 0 {
 		last := s.stack[len(s.stack)-1]
-		s.sample(last, 0)
+		s.sample(last)
 	} else {
 		// First time sample.
 		s.last = time.Now()
@@ -81,7 +89,7 @@ func (s *Stats) Done() {
 		return
 	}
 	last := s.stack[len(s.stack)-1]
-	s.sample(last, 1)
+	s.sample(last)
 	s.stack = s.stack[:len(s.stack)-1]
 }