github.com/hellobchain/third_party@v0.0.0-20230331131523-deb0478a2e52/felixge/httpsnoop/capture_metrics.go (about)

     1  package httpsnoop
     2  
     3  import (
     4  	"github.com/hellobchain/newcryptosm/http"
     5  	"io"
     6  	"sync"
     7  	"time"
     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  		lock          sync.Mutex
    44  		hooks         = Hooks{
    45  			WriteHeader: func(next WriteHeaderFunc) WriteHeaderFunc {
    46  				return func(code int) {
    47  					next(code)
    48  					lock.Lock()
    49  					defer lock.Unlock()
    50  					if !headerWritten {
    51  						m.Code = code
    52  						headerWritten = true
    53  					}
    54  				}
    55  			},
    56  
    57  			Write: func(next WriteFunc) WriteFunc {
    58  				return func(p []byte) (int, error) {
    59  					n, err := next(p)
    60  					lock.Lock()
    61  					defer lock.Unlock()
    62  					m.Written += int64(n)
    63  					headerWritten = true
    64  					return n, err
    65  				}
    66  			},
    67  
    68  			ReadFrom: func(next ReadFromFunc) ReadFromFunc {
    69  				return func(src io.Reader) (int64, error) {
    70  					n, err := next(src)
    71  					lock.Lock()
    72  					defer lock.Unlock()
    73  					headerWritten = true
    74  					m.Written += n
    75  					return n, err
    76  				}
    77  			},
    78  		}
    79  	)
    80  
    81  	fn(Wrap(w, hooks))
    82  	m.Duration = time.Since(start)
    83  	return m
    84  }