capture_metrics.go raw

   1  package httpsnoop
   2  
   3  import (
   4  	"io"
   5  	"net/http"
   6  	"time"
   7  )
   8  
   9  // Metrics holds metrics captured from CaptureMetrics.
  10  type Metrics struct {
  11  	// Code is the first http response code passed to the WriteHeader func of
  12  	// the ResponseWriter. If no such call is made, a default code of 200 is
  13  	// assumed instead.
  14  	Code int
  15  	// Duration is the time it took to execute the handler.
  16  	Duration time.Duration
  17  	// Written is the number of bytes successfully written by the Write or
  18  	// ReadFrom function of the ResponseWriter. ResponseWriters may also write
  19  	// data to their underlaying connection directly (e.g. headers), but those
  20  	// are not tracked. Therefor the number of Written bytes will usually match
  21  	// the size of the response body.
  22  	Written int64
  23  }
  24  
  25  // CaptureMetrics wraps the given hnd, executes it with the given w and r, and
  26  // returns the metrics it captured from it.
  27  func CaptureMetrics(hnd http.Handler, w http.ResponseWriter, r *http.Request) Metrics {
  28  	return CaptureMetricsFn(w, func(ww http.ResponseWriter) {
  29  		hnd.ServeHTTP(ww, r)
  30  	})
  31  }
  32  
  33  // CaptureMetricsFn wraps w and calls fn with the wrapped w and returns the
  34  // resulting metrics. This is very similar to CaptureMetrics (which is just
  35  // sugar on top of this func), but is a more usable interface if your
  36  // application doesn't use the Go http.Handler interface.
  37  func CaptureMetricsFn(w http.ResponseWriter, fn func(http.ResponseWriter)) Metrics {
  38  	m := Metrics{Code: http.StatusOK}
  39  	m.CaptureMetrics(w, fn)
  40  	return m
  41  }
  42  
  43  // CaptureMetrics wraps w and calls fn with the wrapped w and updates
  44  // Metrics m with the resulting metrics. This is similar to CaptureMetricsFn,
  45  // but allows one to customize starting Metrics object.
  46  func (m *Metrics) CaptureMetrics(w http.ResponseWriter, fn func(http.ResponseWriter)) {
  47  	var (
  48  		start         = time.Now()
  49  		headerWritten bool
  50  		hooks         = Hooks{
  51  			WriteHeader: func(next WriteHeaderFunc) WriteHeaderFunc {
  52  				return func(code int) {
  53  					next(code)
  54  
  55  					if !(code >= 100 && code <= 199) && !headerWritten {
  56  						m.Code = code
  57  						headerWritten = true
  58  					}
  59  				}
  60  			},
  61  
  62  			Write: func(next WriteFunc) WriteFunc {
  63  				return func(p []byte) (int, error) {
  64  					n, err := next(p)
  65  
  66  					m.Written += int64(n)
  67  					headerWritten = true
  68  					return n, err
  69  				}
  70  			},
  71  
  72  			ReadFrom: func(next ReadFromFunc) ReadFromFunc {
  73  				return func(src io.Reader) (int64, error) {
  74  					n, err := next(src)
  75  
  76  					headerWritten = true
  77  					m.Written += n
  78  					return n, err
  79  				}
  80  			},
  81  		}
  82  	)
  83  
  84  	fn(Wrap(w, hooks))
  85  	m.Duration += time.Since(start)
  86  }
  87