github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/httpsnoop/capture_metrics.go (about)

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