internal/middleware/requests.go

// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package middleware

import (
	"context"
	"net/http"
	"sort"
	"sync"

	"golang.org/x/pkgsite/internal"
)

var (
	requestMapMu sync.Mutex
	requestMap   = map[string]*internal.RequestInfo{}
)

// RequestInfo adds information about the request to a context.
// It also stores it while the request is active.
// [ActiveRequests] retrieves all stored requests.
func RequestInfo() Middleware {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ri := internal.NewRequestInfo(r)
			ctx, cancel := context.WithCancelCause(r.Context())
			ri.Cancel = cancel

			// If the request has a trace ID, store it in the requestMap while
			// it is active.
			if ri.TraceID != "" {
				requestMapMu.Lock()
				requestMap[ri.TraceID] = ri
				requestMapMu.Unlock()

				defer func() {
					requestMapMu.Lock()
					delete(requestMap, ri.TraceID)
					requestMapMu.Unlock()
				}()
			}

			ctx = internal.NewContextWithRequestInfo(ctx, ri)
			h.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}

// ActiveRequests returns all requests that are currently being handled by the server,
// sorted by start time.
func ActiveRequests() []*internal.RequestInfo {
	requestMapMu.Lock()
	defer requestMapMu.Unlock()
	var ris []*internal.RequestInfo
	for _, ri := range requestMap {
		ris = append(ris, ri)
	}
	sort.Slice(ris, func(i, j int) bool { return ris[i].Start.Before(ris[j].Start) })
	return ris
}

// RequestForTraceID returns the active request with the given trace ID,
// or nil if there is no such request.
func RequestForTraceID(id string) *internal.RequestInfo {
	requestMapMu.Lock()
	defer requestMapMu.Unlock()
	return requestMap[id]
}