github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/querier/stats/time_middleware.go (about)

     1  package stats
     2  
     3  import (
     4  	"net/http"
     5  	"time"
     6  )
     7  
     8  // WallTimeMiddleware tracks the wall time.
     9  type WallTimeMiddleware struct{}
    10  
    11  // NewWallTimeMiddleware makes a new WallTimeMiddleware.
    12  func NewWallTimeMiddleware() WallTimeMiddleware {
    13  	return WallTimeMiddleware{}
    14  }
    15  
    16  // Wrap implements middleware.Interface.
    17  func (m WallTimeMiddleware) Wrap(next http.Handler) http.Handler {
    18  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    19  		if !IsEnabled(r.Context()) {
    20  			next.ServeHTTP(w, r)
    21  			return
    22  		}
    23  
    24  		startTime := time.Now()
    25  		next.ServeHTTP(w, r)
    26  
    27  		stats := FromContext(r.Context())
    28  		stats.AddWallTime(time.Since(startTime))
    29  	})
    30  }