зеркало из https://github.com/golang/pkgsite.git
74 строки
2.1 KiB
Go
74 строки
2.1 KiB
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 internal
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// RequestInfo is information about an HTTP request.
|
|
type RequestInfo struct {
|
|
Request *http.Request
|
|
TrimmedURL string // URL with the scheme and host removed
|
|
TraceID string // extracted from request header
|
|
Start time.Time // when the request began
|
|
State atomic.Value // string describing current state; see [RequestState]
|
|
Cancel func(error) // function that cancels the request's context
|
|
}
|
|
|
|
func NewRequestInfo(r *http.Request) *RequestInfo {
|
|
turl := *r.URL
|
|
turl.Scheme = ""
|
|
turl.Host = ""
|
|
turl.User = nil
|
|
ri := &RequestInfo{
|
|
Request: r,
|
|
TrimmedURL: turl.String(),
|
|
TraceID: r.Header.Get("X-Cloud-Trace-Context"),
|
|
Start: time.Now(),
|
|
}
|
|
ri.State.Store("")
|
|
return ri
|
|
}
|
|
|
|
// requestInfoKey is the type of the context key for RequestInfos.
|
|
type requestInfoKey struct{}
|
|
|
|
// NewContextWithRequestInfo creates a new context from ctx that adds the trace ID.
|
|
func NewContextWithRequestInfo(ctx context.Context, ri *RequestInfo) context.Context {
|
|
return context.WithValue(ctx, requestInfoKey{}, ri)
|
|
}
|
|
|
|
// RequestInfoFromContext retrieves the trace ID from the context, or a zero one
|
|
// if it isn't there.
|
|
func RequestInfoFromContext(ctx context.Context) *RequestInfo {
|
|
ri, _ := ctx.Value(requestInfoKey{}).(*RequestInfo)
|
|
if ri == nil {
|
|
return &RequestInfo{}
|
|
}
|
|
return ri
|
|
}
|
|
|
|
// RequestState sets the State field of the current request.
|
|
// It returns a function that restores the value of the field.
|
|
// Typical use:
|
|
//
|
|
// defer RequestState(ctx, "frobbing")()
|
|
func RequestState(ctx context.Context, s string) func() {
|
|
ri := RequestInfoFromContext(ctx)
|
|
old := ri.State.Load()
|
|
if old == nil {
|
|
// This RequestInfo was not initialized with NewRequestInfo.
|
|
// It is the zero value returned by RequestInfoFromContext
|
|
// when there is no RequestInfo in the context. Do nothing in this case.
|
|
return func() {}
|
|
}
|
|
ri.State.Store(s)
|
|
return func() { ri.State.Store(old) }
|
|
}
|