Make flattening work for MemStats

stats/mem.go

@@ -4,10 +4,7 @@
 
 package stats
 
-import (
-	"encoding/json"
-	"runtime"
-)
+import "runtime"
 
 // BasicMemStats includes a few of the fields from runtime.MemStats suitable for
 // general logging.
@@ -32,52 +29,48 @@ type BasicMemStats struct {
 	PauseTotalNs uint64
 }
 
+type MemStatsPlaceholder interface{}
+
 // MemStatsWrapper wraps runtime.MemStats with an optionally less verbose JSON
 // representation. The JSON field names correspond exactly to the runtime field
 // names to avoid reimplementing the entire struct.
 type MemStatsWrapper struct {
-	basic   *BasicMemStats
-	full    *runtime.MemStats
-	verbose bool
+	MemStatsPlaceholder `json:"Memory"`
+
+	basic *BasicMemStats
+	cache *runtime.MemStats
 }
 
 func NewMemStatsWrapper(verbose bool) *MemStatsWrapper {
-	stats := &MemStatsWrapper{
-		verbose: verbose,
-		full:    &runtime.MemStats{},
-	}
-	if !verbose {
+	stats := &MemStatsWrapper{cache: &runtime.MemStats{}}
+
+	if verbose {
+		stats.MemStatsPlaceholder = stats.cache
+	} else {
 		stats.basic = &BasicMemStats{}
+		stats.MemStatsPlaceholder = stats.basic
 	}
 	return stats
 }
 
-func (s *MemStatsWrapper) MarshalJSON() ([]byte, error) {
-	if s.verbose {
-		return json.Marshal(s.full)
-	} else {
-		return json.Marshal(s.basic)
-	}
-}
-
 // Update fetches the current memstats from runtime and resets the cache.
 func (s *MemStatsWrapper) Update() {
-	runtime.ReadMemStats(s.full)
+	runtime.ReadMemStats(s.cache)
 
-	if !s.verbose {
+	if s.basic != nil {
 		// Gross, but any decent editor can generate this in a couple commands.
-		s.basic.Alloc = s.full.Alloc
-		s.basic.TotalAlloc = s.full.TotalAlloc
-		s.basic.Sys = s.full.Sys
-		s.basic.Lookups = s.full.Lookups
-		s.basic.Mallocs = s.full.Mallocs
-		s.basic.Frees = s.full.Frees
-		s.basic.HeapAlloc = s.full.HeapAlloc
-		s.basic.HeapSys = s.full.HeapSys
-		s.basic.HeapIdle = s.full.HeapIdle
-		s.basic.HeapInuse = s.full.HeapInuse
-		s.basic.HeapReleased = s.full.HeapReleased
-		s.basic.HeapObjects = s.full.HeapObjects
-		s.basic.PauseTotalNs = s.full.PauseTotalNs
+		s.basic.Alloc = s.cache.Alloc
+		s.basic.TotalAlloc = s.cache.TotalAlloc
+		s.basic.Sys = s.cache.Sys
+		s.basic.Lookups = s.cache.Lookups
+		s.basic.Mallocs = s.cache.Mallocs
+		s.basic.Frees = s.cache.Frees
+		s.basic.HeapAlloc = s.cache.HeapAlloc
+		s.basic.HeapSys = s.cache.HeapSys
+		s.basic.HeapIdle = s.cache.HeapIdle
+		s.basic.HeapInuse = s.cache.HeapInuse
+		s.basic.HeapReleased = s.cache.HeapReleased
+		s.basic.HeapObjects = s.cache.HeapObjects
+		s.basic.PauseTotalNs = s.cache.PauseTotalNs
 	}
 }

stats/stats.go

@@ -85,7 +85,7 @@ type Stats struct {
 	IPv4Peers PeerStats `json:"Peers.IPv4"`
 	IPv6Peers PeerStats `json:"Peers.IPv6"`
 
-	MemStats *MemStatsWrapper `json:"Memory,omitempty"`
+	*MemStatsWrapper `json:",omitempty"`
 
 	events             chan int
 	ipv4PeerEvents     chan int
@@ -113,7 +113,7 @@ func New(cfg config.StatsConfig) *Stats {
 	}
 
 	if cfg.IncludeMem {
-		s.MemStats = NewMemStatsWrapper(cfg.VerboseMem)
+		s.MemStatsWrapper = NewMemStatsWrapper(cfg.VerboseMem)
 		s.recordMemStats = time.NewTicker(cfg.MemUpdateInterval.Duration).C
 	}
 
@@ -174,7 +174,7 @@ func (s *Stats) handleEvents() {
 			s.ResponseTime.P95.AddSample(f)
 
 		case <-s.recordMemStats:
-			s.MemStats.Update()
+			s.MemStatsWrapper.Update()
 		}
 	}
 }

stats/struct_flattener.go

@@ -11,7 +11,7 @@ func isEmptyValue(v reflect.Value) bool {
 	return v.Interface() == reflect.Zero(v.Type()).Interface()
 }
 
-func keyForField(field reflect.StructField, v reflect.Value) string {
+func keyForField(field reflect.StructField, v reflect.Value) (string, bool) {
 	if tag := field.Tag.Get("json"); tag != "" {
 		tokens := strings.SplitN(tag, ",", 2)
 		name := tokens[0]
@@ -22,13 +22,29 @@ func keyForField(field reflect.StructField, v reflect.Value) string {
 		}
 
 		if name == "-" || strings.Contains(opts, "omitempty") && isEmptyValue(v) {
-			return ""
+			return "", false
 		} else if name != "" {
-			return name
+			return name, false
 		}
 	}
 
-	return field.Name
+	if field.Anonymous {
+		return "", true
+	}
+	return field.Name, false
+}
+
+func extractValue(val, fallback reflect.Value) reflect.Value {
+	switch val.Kind() {
+	case reflect.Struct:
+		return val
+	case reflect.Ptr:
+		return extractValue(val.Elem(), fallback)
+	case reflect.Interface:
+		return extractValue(val.Elem(), fallback)
+	default:
+		return fallback
+	}
 }
 
 func recursiveFlatten(val reflect.Value, prefix string, output FlatMap) int {
@@ -38,17 +54,28 @@ func recursiveFlatten(val reflect.Value, prefix string, output FlatMap) int {
 	for i := 0; i < val.NumField(); i++ {
 		child := val.Field(i)
 		childType := valType.Field(i)
-		key := prefix + keyForField(childType, child)
+		childPrefix := ""
 
-		if childType.PkgPath != "" || key == "" {
+		key, anonymous := keyForField(childType, child)
+
+		if childType.PkgPath != "" || (key == "" && !anonymous) {
 			continue
-		} else if child.Kind() == reflect.Struct {
-			if recursiveFlatten(child, key+".", output) != 0 {
+		}
+
+		child = extractValue(child, child)
+		if !anonymous {
+			childPrefix = prefix + key + "."
+		}
+
+		if child.Kind() == reflect.Struct {
+			childAdded := recursiveFlatten(child, childPrefix, output)
+			if childAdded != 0 {
+				added += childAdded
 				continue
 			}
 		}
 
-		output[key] = child.Addr().Interface()
+		output[prefix+key] = child.Addr().Interface()
 		added++
 	}