res_tracker/res_tracker.go

package res_tracker

import (
	"fmt"
	"net/http"
	"os"
	"reflect"
	"regexp"
	"runtime/debug"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const maxStackLen = 1000

const OFF = 0
const LITE = 1
const FULL = 2

var mode int32 = OFF
var res2track = &sync.Map{}
var customFormatter = func(item *TrackReportItem) string { return "," }

type TrackedInfo struct {
	Time  time.Time
	Stack []byte
}

type TrackReportItem struct {
	Obj      interface{}
	Duration time.Duration
	Stack    string
}

func InitHandlers() {
	http.HandleFunc("/res_tracker/on", onHandler)
	http.HandleFunc("/res_tracker/off", offHandler)
	http.HandleFunc("/res_tracker/report", reportHandler)
	http.HandleFunc("/res_tracker/reset", resetHandler)
	http.HandleFunc("/res_tracker/status", statusHandler)
}

func SetCustomFormatter(f func(item *TrackReportItem) string) {
	customFormatter = f
}

func onHandler(w http.ResponseWriter, r *http.Request) {
	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(OFF)
	fmt.Fprintf(w, "OFF\n")
}

func reportHandler(w http.ResponseWriter, r *http.Request) {
	for _, line := range ReportLines() {
		fmt.Fprintf(w, line+"\n")
	}
}

func resetHandler(w http.ResponseWriter, r *http.Request) {
	Reset()
	fmt.Fprintf(w, "Reset\n")
}

func statusHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "mode: %v, entries: %d\n", atomic.LoadInt32(&mode), Count())
}

func Enable(m int32) {
	atomic.StoreInt32(&mode, m)
}

func Reset() {
	res2track.Range(func(k, v interface{}) bool {
		res2track.Delete(k)
		return true
	})
}

func IsOn() bool {
	return atomic.LoadInt32(&mode) != OFF
}

func Count() int {
	var i int
	res2track.Range(func(k, v interface{}) bool {
		i++
		return true
	})
	return i
}

func Track(o interface{}) {
	m := atomic.LoadInt32(&mode)
	if m == OFF {
		return
	}
	var stack []byte = nil
	if m == FULL {
		stack = debug.Stack()
	}
	res2track.Store(o, TrackedInfo{time.Now(), stack})
}

func Untrack(o interface{}) {
	if atomic.LoadInt32(&mode) == OFF {
		return
	}
	res2track.Delete(o)
}

func FormatReportItems() []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]

		info, _ := v.(TrackedInfo)
		stackStr := string(info.Stack)

		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] + "..."
			}
		}

		items = append(items, TrackReportItem{Obj: k, Duration: time.Now().Sub(info.Time), Stack: stackStr})
	}

	//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
}

func ReportLines() []string {
	var lines []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 {
		lines = append(lines, fmt.Sprintf("#%d Tracked (%v %p%s duration %v) : %s", idx+1, reflect.TypeOf(item.Obj), item.Obj, customFormatter(&item), item.Duration, item.Stack))
	}
	return lines
}

func ReportToFile(file string) error {
	f, err := os.OpenFile(file, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	for _, line := range ReportLines() {
		if _, err := f.WriteString(line + "\n"); err != nil {
			return err
		}
	}
	return nil
}

func ReportToTmpFile() (string, error) {
	file := fmt.Sprintf("/tmp/track_report.%d", os.Getpid())
	return file, ReportToFile(file)
}