github.com/fjl/memsize@v0.0.2/memsizeui/ui.go (about)

     1  package memsizeui
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"html/template"
     7  	"net/http"
     8  	"reflect"
     9  	"sort"
    10  	"strings"
    11  	"sync"
    12  	"time"
    13  
    14  	"github.com/fjl/memsize"
    15  )
    16  
    17  type Handler struct {
    18  	init     sync.Once
    19  	mux      http.ServeMux
    20  	mu       sync.Mutex
    21  	reports  map[int]Report
    22  	roots    map[string]interface{}
    23  	reportID int
    24  }
    25  
    26  type Report struct {
    27  	ID       int
    28  	Date     time.Time
    29  	Duration time.Duration
    30  	RootName string
    31  	Sizes    memsize.Sizes
    32  }
    33  
    34  type templateInfo struct {
    35  	Roots     []string
    36  	Reports   map[int]Report
    37  	PathDepth int
    38  	Data      interface{}
    39  }
    40  
    41  func (ti *templateInfo) Link(path ...string) string {
    42  	prefix := strings.Repeat("../", ti.PathDepth)
    43  	return prefix + strings.Join(path, "")
    44  }
    45  
    46  func (h *Handler) Add(name string, v interface{}) {
    47  	rv := reflect.ValueOf(v)
    48  	if rv.Kind() != reflect.Ptr || rv.IsNil() {
    49  		panic("root must be non-nil pointer")
    50  	}
    51  	h.mu.Lock()
    52  	if h.roots == nil {
    53  		h.roots = make(map[string]interface{})
    54  	}
    55  	h.roots[name] = v
    56  	h.mu.Unlock()
    57  }
    58  
    59  func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    60  	h.init.Do(func() {
    61  		h.reports = make(map[int]Report)
    62  		h.mux.HandleFunc("/", h.handleRoot)
    63  		h.mux.HandleFunc("/scan", h.handleScan)
    64  		h.mux.HandleFunc("/report/", h.handleReport)
    65  	})
    66  	h.mux.ServeHTTP(w, r)
    67  }
    68  
    69  func (h *Handler) templateInfo(r *http.Request, data interface{}) *templateInfo {
    70  	h.mu.Lock()
    71  	roots := make([]string, 0, len(h.roots))
    72  	for name := range h.roots {
    73  		roots = append(roots, name)
    74  	}
    75  	h.mu.Unlock()
    76  	sort.Strings(roots)
    77  
    78  	return &templateInfo{
    79  		Roots:     roots,
    80  		Reports:   h.reports,
    81  		PathDepth: strings.Count(r.URL.Path, "/") - 1,
    82  		Data:      data,
    83  	}
    84  }
    85  
    86  func (h *Handler) handleRoot(w http.ResponseWriter, r *http.Request) {
    87  	if r.URL.Path != "/" {
    88  		http.NotFound(w, r)
    89  		return
    90  	}
    91  	serveHTML(w, rootTemplate, http.StatusOK, h.templateInfo(r, nil))
    92  }
    93  
    94  func (h *Handler) handleScan(w http.ResponseWriter, r *http.Request) {
    95  	if r.Method != http.MethodPost {
    96  		http.Error(w, "invalid HTTP method, want POST", http.StatusMethodNotAllowed)
    97  		return
    98  	}
    99  	ti := h.templateInfo(r, "Unknown root")
   100  	id, ok := h.scan(r.URL.Query().Get("root"))
   101  	if !ok {
   102  		serveHTML(w, notFoundTemplate, http.StatusNotFound, ti)
   103  		return
   104  	}
   105  	w.Header().Add("Location", ti.Link(fmt.Sprintf("report/%d", id)))
   106  	w.WriteHeader(http.StatusSeeOther)
   107  }
   108  
   109  func (h *Handler) handleReport(w http.ResponseWriter, r *http.Request) {
   110  	var id int
   111  	fmt.Sscan(strings.TrimPrefix(r.URL.Path, "/report/"), &id)
   112  	h.mu.Lock()
   113  	report, ok := h.reports[id]
   114  	h.mu.Unlock()
   115  
   116  	if !ok {
   117  		serveHTML(w, notFoundTemplate, http.StatusNotFound, h.templateInfo(r, "Report not found"))
   118  	} else {
   119  		serveHTML(w, reportTemplate, http.StatusOK, h.templateInfo(r, report))
   120  	}
   121  }
   122  
   123  func (h *Handler) scan(root string) (int, bool) {
   124  	h.mu.Lock()
   125  	defer h.mu.Unlock()
   126  
   127  	val, ok := h.roots[root]
   128  	if !ok {
   129  		return 0, false
   130  	}
   131  	id := h.reportID
   132  	start := time.Now()
   133  	sizes := memsize.Scan(val)
   134  	h.reports[id] = Report{
   135  		ID:       id,
   136  		RootName: root,
   137  		Date:     start.Truncate(1 * time.Second),
   138  		Duration: time.Since(start),
   139  		Sizes:    sizes,
   140  	}
   141  	h.reportID++
   142  	return id, true
   143  }
   144  
   145  func serveHTML(w http.ResponseWriter, tpl *template.Template, status int, ti *templateInfo) {
   146  	w.Header().Set("content-type", "text/html")
   147  	var buf bytes.Buffer
   148  	if err := tpl.Execute(&buf, ti); err != nil {
   149  		http.Error(w, err.Error(), http.StatusInternalServerError)
   150  		return
   151  	}
   152  	buf.WriteTo(w)
   153  }