From cbd23c59033bb3ad66efabe5703e3552794d0624 Mon Sep 17 00:00:00 2001
From: Pavel Shevaev <pacha.shevaev@gmail.com>
Date: Wed, 19 Oct 2022 20:00:22 +0300
Subject: [PATCH] Adding lite and full tracking modes; Using atomic variables;
 Resetting the map only by explicit command using safe operation

---
 res_tracker.go | 110 +++++++++++++++++++++++++++++++------------------
 1 file changed, 70 insertions(+), 40 deletions(-)

diff --git a/res_tracker.go b/res_tracker.go
index c10268b..bf2bb77 100644
--- a/res_tracker.go
+++ b/res_tracker.go
@@ -7,15 +7,20 @@ import (
 	"reflect"
 	"regexp"
 	"runtime/debug"
+	"sort"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
 const maxStackLen = 1000
 
-//NOTE: for simplicity not making it an atomic variable
-var enabled bool = false
+const OFF = 0
+const LITE = 1
+const FULL = 2
+
+var mode int32 = OFF
 var res2track = &sync.Map{}
 
 type TrackedInfo struct {
@@ -38,12 +43,19 @@ func InitHandlers() {
 }
 
 func onHandler(w http.ResponseWriter, r *http.Request) {
-	Enable(true)
-	fmt.Fprintf(w, "ON\n")
+	query := r.URL.Query()
+	m := query.Get("m")
+	if len(m) == 0 || m == "lite" {
+		Enable(LITE)
+		fmt.Fprintf(w, "ON(lite)\n")
+	} else if m == "full" {
+		Enable(FULL)
+		fmt.Fprintf(w, "ON(full)\n")
+	}
 }
 
 func offHandler(w http.ResponseWriter, r *http.Request) {
-	Enable(false)
+	Enable(OFF)
 	fmt.Fprintf(w, "OFF\n")
 }
 
@@ -59,31 +71,27 @@ func resetHandler(w http.ResponseWriter, r *http.Request) {
 }
 
 func statusHandler(w http.ResponseWriter, r *http.Request) {
-	fmt.Fprintf(w, "enabled: %v, entries: %d\n", enabled, Count())
+	fmt.Fprintf(w, "mode: %v, entries: %d\n", atomic.LoadInt32(&mode), Count())
 }
 
-func Enable(flag bool) {
-	enabled = flag
-	if enabled {
-		Reset()
-	}
+func Enable(m int32) {
+	atomic.StoreInt32(&mode, m)
 }
 
 func Reset() {
-	res2track = &sync.Map{}
+	res2track.Range(func(k, v interface{}) bool {
+		res2track.Delete(k)
+		return true
+	})
 }
 
 func IsOn() bool {
-	return enabled
+	return atomic.LoadInt32(&mode) != OFF
 }
 
 func Count() int {
-	return count(res2track)
-}
-
-func count(m *sync.Map) int {
 	var i int
-	m.Range(func(k, v interface{}) bool {
+	res2track.Range(func(k, v interface{}) bool {
 		i++
 		return true
 	})
@@ -91,55 +99,75 @@ func count(m *sync.Map) int {
 }
 
 func Track(o interface{}) {
-	if !enabled {
+	m := atomic.LoadInt32(&mode)
+	if m == OFF {
 		return
 	}
-	stack := debug.Stack()
+	var stack []byte = nil
+	if m == FULL {
+		stack = debug.Stack()
+	}
 	res2track.Store(o, TrackedInfo{time.Now(), stack})
 }
 
 func Untrack(o interface{}) {
-	if !enabled {
+	if atomic.LoadInt32(&mode) == OFF {
 		return
 	}
 	res2track.Delete(o)
 }
 
 func FormatReportItems() []TrackReportItem {
-	var items []TrackReportItem
-
 	parensReg := regexp.MustCompile(`\([^\)]+\)`)
 	fileReg := regexp.MustCompile(`at\s+.*?/([^/]+\.go:\d+)`)
 
+	//1. let's get a quick snapshot of the map
+	var snapshot []interface{}
 	res2track.Range(func(k, v interface{}) bool {
+		snapshot = append(snapshot, k)
+		snapshot = append(snapshot, v)
+		return true
+	})
+
+	//2. let's add formatted items
+	var items []TrackReportItem
+	for i := 0; i < len(snapshot); i += 2 {
+		k := snapshot[i]
+		v := snapshot[i+1]
 
 		t := reflect.TypeOf(k)
 		info, _ := v.(TrackedInfo)
 		stackStr := string(info.Stack)
 
-		//let's remove first stack entries
-		for i := 0; i < 5; i++ {
-			newLineIdx := strings.Index(stackStr, "\n")
-			if newLineIdx != -1 {
-				stackStr = stackStr[newLineIdx+1:]
-			} else {
-				break
+		if len(stackStr) > 0 {
+			//let's remove first stack entries
+			for i := 0; i < 5; i++ {
+				newLineIdx := strings.Index(stackStr, "\n")
+				if newLineIdx != -1 {
+					stackStr = stackStr[newLineIdx+1:]
+				} else {
+					break
+				}
+			}
+
+			stackStr = strings.ReplaceAll(stackStr, "  ", " ")
+			stackStr = strings.ReplaceAll(stackStr, "\n", " < ")
+			stackStr = strings.ReplaceAll(stackStr, ") <", ") at")
+			stackStr = parensReg.ReplaceAllString(stackStr, "(..)")
+			stackStr = fileReg.ReplaceAllString(stackStr, "../$1")
+			if len(stackStr) > maxStackLen {
+				stackStr = stackStr[:maxStackLen] + "..."
 			}
 		}
 
-		stackStr = strings.ReplaceAll(stackStr, "  ", " ")
-		stackStr = strings.ReplaceAll(stackStr, "\n", " < ")
-		stackStr = strings.ReplaceAll(stackStr, ") <", ") at")
-		stackStr = parensReg.ReplaceAllString(stackStr, "(..)")
-		stackStr = fileReg.ReplaceAllString(stackStr, "../$1")
-		if len(stackStr) > maxStackLen {
-			stackStr = stackStr[:maxStackLen] + "..."
-		}
-
 		items = append(items, TrackReportItem{T: t, Duration: time.Now().Sub(info.Time), Stack: stackStr})
+	}
 
-		return true
+	//3. let's sort items by duration descending
+	sort.Slice(items, func(i, j int) bool {
+		return items[i].Duration > items[j].Duration
 	})
+
 	return items
 }
 
@@ -149,8 +177,10 @@ func ReportLines() []string {
 	now := time.Now()
 
 	items := FormatReportItems()
+
 	lines = append(lines, fmt.Sprintf("=== Tracked connections report %s (total: %d) ===", now.Format("01-02-2006 15:04:05"), len(items)))
 	for idx, item := range items {
+		//TODO: obtain somehow an actual pointer of the object or provide a hook for a custom formatter
 		lines = append(lines, fmt.Sprintf("#%d Tracked (%v, duration %v) : %s", idx+1, item.T, item.Duration, item.Stack))
 	}
 	return lines