res_tracker/res_tracker.go

package res_tracker

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

const maxStackLen = 1000

//NOTE: for simplicity not making it an atomic variable
var enabled bool = false
var res2track = &sync.Map{}

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

type TrackReportItem struct {
	T        reflect.Type
	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 onHandler(w http.ResponseWriter, r *http.Request) {
	Enable(true)
	fmt.Fprintf(w, "ON\n")
}

func offHandler(w http.ResponseWriter, r *http.Request) {
	Enable(false)
	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, "enabled: %v, entries: %d\n", enabled, Count())
}

func Enable(flag bool) {
	enabled = flag
	if enabled {
		Reset()
	}
}

func Reset() {
	res2track = &sync.Map{}
}

func IsOn() bool {
	return enabled
}

func Count() int {
	return count(res2track)
}

func count(m *sync.Map) int {
	var i int
	m.Range(func(k, v interface{}) bool {
		i++
		return true
	})
	return i
}

func Track(o interface{}) {
	if !enabled {
		return
	}
	stack := debug.Stack()
	res2track.Store(o, TrackedInfo{time.Now(), stack})
}

func Untrack(o interface{}) {
	if !enabled {
		return
	}
	res2track.Delete(o)
}

func FormatReportItems() []TrackReportItem {
	var items []TrackReportItem

	parensReg := regexp.MustCompile(`\([^\)]+\)`)
	fileReg := regexp.MustCompile(`at\s+.*?/([^/]+\.go:\d+)`)

	res2track.Range(func(k, v interface{}) bool {

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

		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
	})
	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, duration %v) : %s", idx+1, item.T, 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)
}