forked from TrueCloudLab/frostfs-s3-gw
108 lines
2.5 KiB
Go
108 lines
2.5 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"sync"
|
|
|
|
"git.frostfs.info/TrueCloudLab/frostfs-observability/tracing"
|
|
"go.opentelemetry.io/otel/attribute"
|
|
semconv "go.opentelemetry.io/otel/semconv/v1.18.0"
|
|
"go.opentelemetry.io/otel/trace"
|
|
)
|
|
|
|
type traceResponseWriter struct {
|
|
sync.Once
|
|
http.ResponseWriter
|
|
|
|
ctx context.Context
|
|
span trace.Span
|
|
}
|
|
|
|
func (lrw *traceResponseWriter) WriteHeader(code int) {
|
|
lrw.Do(func() {
|
|
lrw.span.SetAttributes(
|
|
semconv.HTTPStatusCode(code),
|
|
)
|
|
|
|
carrier := &httpResponseCarrier{resp: lrw.ResponseWriter}
|
|
tracing.Propagator.Inject(lrw.ctx, carrier)
|
|
|
|
lrw.ResponseWriter.WriteHeader(code)
|
|
lrw.span.End()
|
|
})
|
|
}
|
|
|
|
func (lrw *traceResponseWriter) Flush() {
|
|
if f, ok := lrw.ResponseWriter.(http.Flusher); ok {
|
|
f.Flush()
|
|
}
|
|
}
|
|
|
|
type httpResponseCarrier struct {
|
|
resp http.ResponseWriter
|
|
}
|
|
|
|
func (h httpResponseCarrier) Get(key string) string {
|
|
return h.resp.Header().Get(key)
|
|
}
|
|
|
|
func (h httpResponseCarrier) Set(key string, value string) {
|
|
h.resp.Header().Set(key, value)
|
|
}
|
|
|
|
func (h httpResponseCarrier) Keys() []string {
|
|
result := make([]string, 0, len(h.resp.Header()))
|
|
for key := range h.resp.Header() {
|
|
result = append(result, key)
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
type httpRequestCarrier struct {
|
|
req *http.Request
|
|
}
|
|
|
|
func (c *httpRequestCarrier) Get(key string) string {
|
|
bytes := c.req.Header.Get(key)
|
|
if len(bytes) == 0 {
|
|
return ""
|
|
}
|
|
return bytes
|
|
}
|
|
|
|
func (c *httpRequestCarrier) Set(key string, value string) {
|
|
c.req.Response.Header.Set(key, value)
|
|
}
|
|
|
|
func (c *httpRequestCarrier) Keys() []string {
|
|
result := make([]string, 0, len(c.req.Header))
|
|
for key := range c.req.Header {
|
|
result = append(result, key)
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
func extractHTTPTraceInfo(ctx context.Context, req *http.Request) context.Context {
|
|
if req == nil {
|
|
return ctx
|
|
}
|
|
carrier := &httpRequestCarrier{req: req}
|
|
return tracing.Propagator.Extract(ctx, carrier)
|
|
}
|
|
|
|
// StartHTTPServerSpan starts root HTTP server span.
|
|
func StartHTTPServerSpan(r *http.Request, operationName string, opts ...trace.SpanStartOption) (context.Context, trace.Span) {
|
|
ctx := extractHTTPTraceInfo(r.Context(), r)
|
|
opts = append(opts, trace.WithAttributes(
|
|
attribute.String("s3.client_address", r.RemoteAddr),
|
|
attribute.String("s3.path", r.Host),
|
|
attribute.String("s3.request_id", GetRequestID(r.Context())),
|
|
semconv.HTTPMethod(r.Method),
|
|
semconv.RPCService("frostfs-s3-gw"),
|
|
attribute.String("s3.query", r.RequestURI),
|
|
), trace.WithSpanKind(trace.SpanKindServer))
|
|
return tracing.StartSpanFromContext(ctx, operationName, opts...)
|
|
}
|