1 files changed, 107 insertions, 0 deletions
diff --git a/internal/httptransport/metered_round_tripper.go b/internal/httptransport/metered_round_tripper.go
new file mode 100644
index 00000000..d04979a2
--- /dev/null
+++ b/internal/httptransport/metered_round_tripper.go
@@ -0,0 +1,107 @@
+package httptransport
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptrace"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	log "github.com/sirupsen/logrus"
+)
+
+// Options to configure a http.Transport
+type Options func(transport *http.Transport)
+
+type meteredRoundTripper struct {
+	next        http.RoundTripper
+	name        string
+	tracer      *prometheus.HistogramVec
+	durations   *prometheus.HistogramVec
+	counter     *prometheus.CounterVec
+	ttfbTimeout time.Duration
+}
+
+// NewMeteredRoundTripper will create a custom http.RoundTripper that can be used with an http.Client.
+// The RoundTripper will report metrics based on the collectors passed.
+func NewMeteredRoundTripper(name string, tracerVec, durationsVec *prometheus.
+	HistogramVec, counterVec *prometheus.CounterVec, ttfbTimeout time.Duration) http.RoundTripper {
+	return &meteredRoundTripper{
+		next:        InternalTransport,
+		name:        name,
+		tracer:      tracerVec,
+		durations:   durationsVec,
+		counter:     counterVec,
+		ttfbTimeout: ttfbTimeout,
+	}
+}
+
+// WithFileProtocol option to be used while ReconfigureMeteredRoundTripper
+func WithFileProtocol(protocol string, rt http.RoundTripper) Options {
+	return func(transport *http.Transport) {
+		transport.RegisterProtocol(protocol, rt)
+	}
+}
+
+// ReconfigureMeteredRoundTripper clones meteredRoundTripper and applies options to the transport
+func ReconfigureMeteredRoundTripper(rt http.RoundTripper, opts ...Options) http.RoundTripper {
+	mrt, ok := rt.(*meteredRoundTripper)
+	if !ok {
+		return nil
+	}
+
+	t := clone(mrt.next.(*http.Transport))
+	for _, opt := range opts {
+		opt(t)
+	}
+
+	mrt.next = t
+
+	return mrt
+}
+
+// RoundTripper wraps the original http.Transport into a meteredRoundTripper which
+// reports metrics on request duration, tracing and request count
+func (mrt *meteredRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	start := time.Now()
+
+	ctx := httptrace.WithClientTrace(r.Context(), mrt.newTracer(start))
+	ctx, cancel := context.WithCancel(ctx)
+
+	timer := time.AfterFunc(mrt.ttfbTimeout, cancel)
+	defer timer.Stop()
+
+	r = r.WithContext(ctx)
+
+	resp, err := mrt.next.RoundTrip(r)
+	if err != nil {
+		mrt.counter.WithLabelValues("error").Inc()
+		return nil, err
+	}
+
+	mrt.logResponse(r, resp)
+
+	statusCode := strconv.Itoa(resp.StatusCode)
+	mrt.durations.WithLabelValues(statusCode).Observe(time.Since(start).Seconds())
+	mrt.counter.WithLabelValues(statusCode).Inc()
+
+	return resp, nil
+}
+
+func (mrt *meteredRoundTripper) logResponse(req *http.Request, resp *http.Response) {
+	if log.GetLevel() == log.TraceLevel {
+		l := log.WithFields(log.Fields{
+			"client_name":     mrt.name,
+			"req_url":         req.URL.String(),
+			"res_status_code": resp.StatusCode,
+		})
+
+		for header, value := range resp.Header {
+			l = l.WithField(strings.ToLower(header), strings.Join(value, ";"))
+		}
+
+		l.Traceln("response")
+	}
+}