github.com/msales/pkg/v3@v3.24.0/httpx/middleware/stats.go (about)

     1  package middleware
     2  
     3  import (
     4  	"net/http"
     5  	"strconv"
     6  
     7  	"github.com/msales/pkg/v3/stats"
     8  )
     9  
    10  // TagsFunc returns a set of tags from a request
    11  type TagsFunc func(*http.Request) []interface{}
    12  
    13  // DefaultTags extracts the method and path from the request.
    14  func DefaultTags(r *http.Request) []interface{} {
    15  	return []interface{}{
    16  		"method", r.Method,
    17  		"path", r.URL.Path,
    18  	}
    19  }
    20  
    21  // WithRequestStats collects statistics about the request.
    22  func WithRequestStats(h http.Handler, fns ...TagsFunc) http.Handler {
    23  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    24  		tags := prepareTags(r, fns)
    25  
    26  		s, ok := stats.FromContext(r.Context())
    27  		if !ok {
    28  			s = stats.Null
    29  		}
    30  
    31  		_ = s.Inc("request.start", 1, 1.0, tags...)
    32  
    33  		rw := NewResponseWriter(w)
    34  		h.ServeHTTP(rw, r)
    35  
    36  		cpltTags := make([]interface{}, len(tags)+2)
    37  		cpltTags[0] = "status"
    38  		cpltTags[1] = strconv.FormatInt(int64(rw.Status()), 10)
    39  		copy(cpltTags[2:], tags)
    40  		_ = s.Inc("request.complete", 1, 1.0, cpltTags...)
    41  	})
    42  }
    43  
    44  // WithResponseTime reports the response time.
    45  func WithResponseTime(h http.Handler, fns ...TagsFunc) http.Handler {
    46  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    47  		tags := prepareTags(r, fns)
    48  		t := stats.Time(r.Context(), "response.time", 1.0, tags...)
    49  		defer t.Done()
    50  
    51  		h.ServeHTTP(w, r)
    52  	})
    53  }
    54  
    55  // prepareTags resolves tags in accordance to provided functions and falls back to defaults in no custom tag functions were provided.
    56  func prepareTags(r *http.Request, fns []TagsFunc) []interface{} {
    57  	if len(fns) == 0 {
    58  		fns = []TagsFunc{DefaultTags}
    59  	}
    60  
    61  	var tags []interface{}
    62  	for _, fn := range fns {
    63  		tags = append(tags, fn(r)...)
    64  	}
    65  
    66  	return tags
    67  }