internal/logging/slowlogs/middleware.go

package slowlogs

import (
	"context"
	"net/http"
	"time"

	"github.com/sirupsen/logrus"
	"gitlab.com/gitlab-org/labkit/log"
)

// LoggerFunc Define a type for the logging function
type LoggerFunc func(r *http.Request) *logrus.Entry

const SlowRequestTimeThreshold = 500 * time.Millisecond

// NewMiddleware returns middleware which logs timing info if downstream time is more than 0.5 seconds
func NewMiddleware(handler http.Handler, logger LoggerFunc, slowRequestThreshold time.Duration) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		// update the request context with timing info
		ctx := withTimingInfo(r.Context())
		r = r.WithContext(ctx)

		handler.ServeHTTP(w, r)

		totalDownstreamTime := time.Since(start)
		if totalDownstreamTime >= slowRequestThreshold {
			logger(r).WithFields(log.Fields{
				"total_execution_time_ms":      totalDownstreamTime.Milliseconds(),
				"middleware_time_distribution": prepareMiddlewareTimingInfoInMS(r.Context()),
				"other_time_distribution":      prepareOtherTimingInfoInMS(r.Context()),
			}).Info("Request time exceeded 0.5 seconds")
		}
	})
}

// prepareMiddlewareTimingInfoLog return timing info as map of string and duration in milliseconds.
func prepareMiddlewareTimingInfoInMS(ctx context.Context) map[string]int64 {
	info := ctx.Value(timingContextKey).(*timingInfo)

	return prepareTimingInfoInMS(info.MiddlewareTimings)
}

// prepareOtherTimingInfoLog return timing info as map of string and duration in milliseconds.
func prepareOtherTimingInfoInMS(ctx context.Context) map[string]int64 {
	info := ctx.Value(timingContextKey).(*timingInfo)

	return prepareTimingInfoInMS(info.OtherTimings)
}

func prepareTimingInfoInMS(timingDurations map[string]time.Duration) map[string]int64 {
	timingsInMillis := make(map[string]int64)
	for key, duration := range timingDurations {
		timingsInMillis[key+"_ms"] = duration.Milliseconds()
	}
	return timingsInMillis
}

// LogHandlerTiming calculates handler execution time excluding downstream handler execution time
func LogHandlerTiming(handler http.Handler, handlerName string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		startLogTime := time.Now()
		defer calculateMiddlewareExecutionTime(r.Context(), startLogTime, handlerName)
		handler.ServeHTTP(w, r)
	})
}

func calculateMiddlewareExecutionTime(ctx context.Context, startTime time.Time, handlerName string) {
	info := getTimingInfo(ctx)
	spentTime := time.Since(startTime)
	info.MiddlewareTimings[handlerName] = spentTime - info.LastDuration
	info.LastDuration = spentTime
}