181 lines · 4.6 KB · Go
RawDownload
package logging import (	"context"	"net/http"	"sync"	"sync/atomic") type endpointKey struct{}type responseStatusKey struct{}type responseHeadersKey struct{}type clientRequestMetadataKey struct{} // ClientRequestMetadata stores immutable downstream request metadata for asynchronous consumers.type ClientRequestMetadata struct {	ClientIP        string	XForwardedFor   string	UserAgent       string	SessionID       string	ParentSessionID string} type responseStatusHolder struct {	status atomic.Int32} type responseHeadersHolder struct {	mu      sync.RWMutex	headers http.Header} func WithEndpoint(ctx context.Context, endpoint string) context.Context {	if ctx == nil {		ctx = context.Background()	}	return context.WithValue(ctx, endpointKey{}, endpoint)} func GetEndpoint(ctx context.Context) string {	if ctx == nil {		return ""	}	if endpoint, ok := ctx.Value(endpointKey{}).(string); ok {		return endpoint	}	return ""} // WithClientRequestMetadata stores a snapshot of downstream request metadata in ctx.func WithClientRequestMetadata(ctx context.Context, metadata ClientRequestMetadata) context.Context {	if ctx == nil {		ctx = context.Background()	}	return context.WithValue(ctx, clientRequestMetadataKey{}, metadata)} // GetClientRequestMetadata returns downstream request metadata stored in ctx.func GetClientRequestMetadata(ctx context.Context) ClientRequestMetadata {	if ctx == nil {		return ClientRequestMetadata{}	}	if metadata, ok := ctx.Value(clientRequestMetadataKey{}).(ClientRequestMetadata); ok {		return metadata	}	return ClientRequestMetadata{}} func WithResponseStatusHolder(ctx context.Context) context.Context {	if ctx == nil {		ctx = context.Background()	}	if holder, ok := ctx.Value(responseStatusKey{}).(*responseStatusHolder); ok && holder != nil {		return ctx	}	return context.WithValue(ctx, responseStatusKey{}, &responseStatusHolder{})} func WithResponseHeadersHolder(ctx context.Context) context.Context {	if ctx == nil {		ctx = context.Background()	}	if holder, ok := ctx.Value(responseHeadersKey{}).(*responseHeadersHolder); ok && holder != nil {		return ctx	}	return context.WithValue(ctx, responseHeadersKey{}, &responseHeadersHolder{})} // WithFreshResponseHeadersHolder starts an isolated upstream response attempt.// Unlike WithResponseHeadersHolder, it always shadows any holder inherited from// the parent request so a later retry cannot observe headers from an earlier// credential or model attempt.func WithFreshResponseHeadersHolder(ctx context.Context) context.Context {	if ctx == nil {		ctx = context.Background()	}	return context.WithValue(ctx, responseHeadersKey{}, &responseHeadersHolder{})} func SetResponseStatus(ctx context.Context, status int) {	if ctx == nil || status <= 0 {		return	}	holder, ok := ctx.Value(responseStatusKey{}).(*responseStatusHolder)	if !ok || holder == nil {		return	}	holder.status.Store(int32(status))} func SetResponseHeaders(ctx context.Context, headers http.Header) {	if ctx == nil {		return	}	holder, ok := ctx.Value(responseHeadersKey{}).(*responseHeadersHolder)	if !ok || holder == nil {		return	}	holder.mu.Lock()	defer holder.mu.Unlock()	holder.headers = cloneHTTPHeader(headers)} // MergeResponseHeaders adds headers observed after the initial HTTP response,// such as quota metadata delivered in a websocket event.func MergeResponseHeaders(ctx context.Context, headers http.Header) {	if ctx == nil || len(headers) == 0 {		return	}	holder, ok := ctx.Value(responseHeadersKey{}).(*responseHeadersHolder)	if !ok || holder == nil {		return	}	holder.mu.Lock()	defer holder.mu.Unlock()	if holder.headers == nil {		holder.headers = make(http.Header, len(headers))	}	for key, values := range headers {		canonicalKey := http.CanonicalHeaderKey(key)		if canonicalKey == "" {			continue		}		holder.headers[canonicalKey] = append([]string(nil), values...)	}} func GetResponseStatus(ctx context.Context) int {	if ctx == nil {		return 0	}	holder, ok := ctx.Value(responseStatusKey{}).(*responseStatusHolder)	if !ok || holder == nil {		return 0	}	return int(holder.status.Load())} func GetResponseHeaders(ctx context.Context) http.Header {	if ctx == nil {		return nil	}	holder, ok := ctx.Value(responseHeadersKey{}).(*responseHeadersHolder)	if !ok || holder == nil {		return nil	}	holder.mu.RLock()	defer holder.mu.RUnlock()	return cloneHTTPHeader(holder.headers)} func cloneHTTPHeader(src http.Header) http.Header {	if len(src) == 0 {		return nil	}	dst := make(http.Header, len(src))	for key, values := range src {		dst[key] = append([]string(nil), values...)	}	return dst}