github.com/rudderlabs/rudder-go-kit@v0.30.0/chiware/stats.go (about)

     1  package chiware
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net/http"
     7  	"strconv"
     8  	"sync/atomic"
     9  	"time"
    10  
    11  	"github.com/go-chi/chi/v5"
    12  
    13  	"github.com/rudderlabs/rudder-go-kit/stats"
    14  )
    15  
    16  type config struct {
    17  	redactUnknownPaths bool
    18  }
    19  
    20  type Option func(*config)
    21  
    22  // RedactUnknownPaths sets the redactUnknownPaths flag.
    23  // If set to true, the path will be redacted if the route is not found.
    24  // If set to false, the path will be used as is.
    25  func RedactUnknownPaths(redactUnknownPaths bool) Option {
    26  	return func(c *config) {
    27  		c.redactUnknownPaths = redactUnknownPaths
    28  	}
    29  }
    30  
    31  func StatMiddleware(ctx context.Context, s stats.Stats, component string, options ...Option) func(http.Handler) http.Handler {
    32  	conf := config{
    33  		redactUnknownPaths: true,
    34  	}
    35  	for _, option := range options {
    36  		option(&conf)
    37  	}
    38  	var concurrentRequests int32
    39  	activeClientCount := s.NewStat(fmt.Sprintf("%s.concurrent_requests_count", component), stats.GaugeType)
    40  	go func() {
    41  		for {
    42  			select {
    43  			case <-ctx.Done():
    44  				return
    45  			case <-time.After(10 * time.Second):
    46  				activeClientCount.Gauge(atomic.LoadInt32(&concurrentRequests))
    47  			}
    48  		}
    49  	}()
    50  
    51  	// getPath retrieves the path from the request.
    52  	// The matched route's template is used if a match is found,
    53  	// otherwise the request's URL path is used instead.
    54  	getPath := func(r *http.Request) string {
    55  		if path := chi.RouteContext(r.Context()).RoutePattern(); path != "" {
    56  			return path
    57  		}
    58  		if conf.redactUnknownPaths {
    59  			return "/redacted"
    60  		}
    61  		return r.URL.Path
    62  	}
    63  	return func(next http.Handler) http.Handler {
    64  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    65  			sw := newStatusCapturingWriter(w)
    66  			start := time.Now()
    67  			atomic.AddInt32(&concurrentRequests, 1)
    68  			defer atomic.AddInt32(&concurrentRequests, -1)
    69  
    70  			next.ServeHTTP(sw, r)
    71  			s.NewSampledTaggedStat(
    72  				fmt.Sprintf("%s.response_time", component),
    73  				stats.TimerType,
    74  				map[string]string{
    75  					"reqType": getPath(r),
    76  					"method":  r.Method,
    77  					"code":    strconv.Itoa(sw.status),
    78  				}).Since(start)
    79  		})
    80  	}
    81  }
    82  
    83  // newStatusCapturingWriter returns a new, properly initialized statusCapturingWriter
    84  func newStatusCapturingWriter(w http.ResponseWriter) *statusCapturingWriter {
    85  	return &statusCapturingWriter{
    86  		ResponseWriter: w,
    87  		status:         http.StatusOK,
    88  	}
    89  }
    90  
    91  // statusCapturingWriter is a response writer decorator that captures the status code.
    92  type statusCapturingWriter struct {
    93  	http.ResponseWriter
    94  	status int
    95  }
    96  
    97  // WriteHeader override the http.ResponseWriter's `WriteHeader` method
    98  func (w *statusCapturingWriter) WriteHeader(status int) {
    99  	w.status = status
   100  	w.ResponseWriter.WriteHeader(status)
   101  }