From 7f9fe4b97058830be58f79e8a54530863d5cd829 Mon Sep 17 00:00:00 2001
From: Jonathan Moody <103143855+moodyjon@users.noreply.github.com>
Date: Mon, 18 Jul 2022 11:52:15 -0400
Subject: [PATCH] [rpc mempool] More tweaks to dynamicMemUsage(). Add
 toggleable assertions for max depth and switch completness. Toggle them when
 running in mempool_test.go. Drop support for reflect.Map, as it's not needed
 at this time.

---
 mempool/mempool_test.go |  6 ++++
 mempool/memusage.go     | 73 +++++++++++++++++++++--------------------
 2 files changed, 43 insertions(+), 36 deletions(-)

diff --git a/mempool/mempool_test.go b/mempool/mempool_test.go
index 17d2d452..1b343cd8 100644
--- a/mempool/mempool_test.go
+++ b/mempool/mempool_test.go
@@ -21,6 +21,12 @@ import (
 	btcutil "github.com/lbryio/lbcutil"
 )
 
+func init() {
+	// Toggle assert & debug messages when running tests.
+	dynamicMemUsageAssert = true
+	dynamicMemUsageDebug = false
+}
+
 // fakeChain is used by the pool harness to provide generated test utxos and
 // a current faked chain height to the pool callbacks.  This, in turn, allows
 // transactions to appear as though they are spending completely valid utxos.
diff --git a/mempool/memusage.go b/mempool/memusage.go
index e65ef8c2..c4030a08 100644
--- a/mempool/memusage.go
+++ b/mempool/memusage.go
@@ -5,83 +5,84 @@
 package mempool
 
 import (
+	"fmt"
 	"reflect"
 )
 
+var (
+	dynamicMemUsageAssert   = false
+	dynamicMemUsageDebug    = false
+	dynamicMemUsageMaxDepth = 10
+)
+
 func dynamicMemUsage(v reflect.Value) uintptr {
-	return _dynamicMemUsage(v, false, 0)
+	return dynamicMemUsageCrawl(v, 0)
 }
 
-func _dynamicMemUsage(v reflect.Value, debug bool, level int) uintptr {
+func dynamicMemUsageCrawl(v reflect.Value, depth int) uintptr {
 	t := v.Type()
 	bytes := t.Size()
-	if debug {
-		println("[", level, "]", t.Kind().String(), "(", t.String(), ") ->", t.Size())
+	if dynamicMemUsageDebug {
+		println("[", depth, "]", t.Kind().String(), "(", t.String(), ") ->", t.Size())
 	}
 
-	// For complex types, we need to peek inside slices/arrays/structs/maps and chase pointers.
+	if depth >= dynamicMemUsageMaxDepth {
+		if dynamicMemUsageAssert {
+			panic("crawl reached maximum depth")
+		}
+		return bytes
+	}
+
+	// For complex types, we need to peek inside slices/arrays/structs and chase pointers.
 	switch t.Kind() {
 	case reflect.Pointer, reflect.Interface:
 		if !v.IsNil() {
-			bytes += _dynamicMemUsage(v.Elem(), debug, level+1)
+			bytes += dynamicMemUsageCrawl(v.Elem(), depth+1)
 		}
 	case reflect.Array, reflect.Slice:
 		for j := 0; j < v.Len(); j++ {
 			vi := v.Index(j)
 			k := vi.Type().Kind()
-			if debug {
-				println("[", level, "] index:", j, "kind:", k.String())
+			if dynamicMemUsageDebug {
+				println("[", depth, "] index:", j, "kind:", k.String())
 			}
-			elemB := uintptr(0)
+			elemBytes := uintptr(0)
 			if t.Kind() == reflect.Array {
 				if (k == reflect.Pointer || k == reflect.Interface) && !vi.IsNil() {
-					elemB += _dynamicMemUsage(vi.Elem(), debug, level+1)
+					elemBytes += dynamicMemUsageCrawl(vi.Elem(), depth+1)
 				}
 			} else { // slice
-				elemB += _dynamicMemUsage(vi, debug, level+1)
+				elemBytes += dynamicMemUsageCrawl(vi, depth+1)
 			}
 			if k == reflect.Uint8 {
 				// short circuit for byte slice/array
-				bytes += elemB * uintptr(v.Len())
-				if debug {
+				bytes += elemBytes * uintptr(v.Len())
+				if dynamicMemUsageDebug {
 					println("...", v.Len(), "elements")
 				}
 				break
 			}
-			bytes += elemB
-		}
-	case reflect.Map:
-		iter := v.MapRange()
-		for iter.Next() {
-			vk := iter.Key()
-			vv := iter.Value()
-			if debug {
-				println("[", level, "] key:", vk.Type().Kind().String())
-			}
-			bytes += _dynamicMemUsage(vk, debug, level+1)
-			if debug {
-				println("[", level, "] value:", vv.Type().Kind().String())
-			}
-			bytes += _dynamicMemUsage(vv, debug, level+1)
-			if debug {
-				println("...", v.Len(), "map elements")
-			}
-			debug = false
+			bytes += elemBytes
 		}
 	case reflect.Struct:
 		for _, f := range reflect.VisibleFields(t) {
 			vf := v.FieldByIndex(f.Index)
 			k := vf.Type().Kind()
-			if debug {
-				println("[", level, "] field:", f.Name, "kind:", k.String())
+			if dynamicMemUsageDebug {
+				println("[", depth, "] field:", f.Name, "kind:", k.String())
 			}
 			if (k == reflect.Pointer || k == reflect.Interface) && !vf.IsNil() {
-				bytes += _dynamicMemUsage(vf.Elem(), debug, level+1)
+				bytes += dynamicMemUsageCrawl(vf.Elem(), depth+1)
 			} else if k == reflect.Array || k == reflect.Slice {
 				bytes -= vf.Type().Size()
-				bytes += _dynamicMemUsage(vf, debug, level+1)
+				bytes += dynamicMemUsageCrawl(vf, depth+1)
 			}
 		}
+	case reflect.Uint8:
+	default:
+		if dynamicMemUsageAssert {
+			panic(fmt.Sprintf("unsupported kind: %v", t.Kind()))
+		}
 	}
 
 	return bytes