|
5 | 5 | "net/http" |
6 | 6 | "strconv" |
7 | 7 | "strings" |
| 8 | + "sync" |
8 | 9 | "time" |
9 | 10 |
|
10 | 11 | "log/slog" |
@@ -130,6 +131,11 @@ func NewWithConfig(logger *slog.Logger, config Config) func(http.Handler) http.H |
130 | 131 | // dump response body |
131 | 132 | bw := newBodyWriter(w, ResponseBodyMaxSize, config.WithResponseBody) |
132 | 133 |
|
| 134 | + // Make sure we create a map only once per request (in case we have multiple middleware instances) |
| 135 | + if v := r.Context().Value(customAttributesCtxKey); v == nil { |
| 136 | + r = r.WithContext(context.WithValue(r.Context(), customAttributesCtxKey, &sync.Map{})) |
| 137 | + } |
| 138 | + |
133 | 139 | defer func() { |
134 | 140 | // Pass thru filters and skip early the code below, to prevent unnecessary processing. |
135 | 141 | for _, filter := range config.Filters { |
@@ -232,9 +238,11 @@ func NewWithConfig(logger *slog.Logger, config Config) func(http.Handler) http.H |
232 | 238 |
|
233 | 239 | // custom context values |
234 | 240 | if v := r.Context().Value(customAttributesCtxKey); v != nil { |
235 | | - switch attrs := v.(type) { |
236 | | - case []slog.Attr: |
237 | | - attributes = append(attributes, attrs...) |
| 241 | + if m, ok := v.(*sync.Map); ok { |
| 242 | + m.Range(func(key, value any) bool { |
| 243 | + attributes = append(attributes, slog.Attr{Key: key.(string), Value: value.(slog.Value)}) |
| 244 | + return true |
| 245 | + }) |
238 | 246 | } |
239 | 247 | } |
240 | 248 |
|
@@ -268,17 +276,17 @@ func GetRequestIDFromContext(ctx context.Context) string { |
268 | 276 | return "" |
269 | 277 | } |
270 | 278 |
|
271 | | -// AddCustomAttributes adds custom attributes to the request context. |
| 279 | +// AddCustomAttributes adds custom attributes to the request context. This func can be called from any handler or middleware, as long as the slog-http middleware is already mounted. |
272 | 280 | func AddCustomAttributes(r *http.Request, attr slog.Attr) { |
273 | | - v := r.Context().Value(customAttributesCtxKey) |
274 | | - if v == nil { |
275 | | - *r = *r.WithContext(context.WithValue(r.Context(), customAttributesCtxKey, []slog.Attr{attr})) |
276 | | - return |
277 | | - } |
| 281 | + AddContextAttributes(r.Context(), attr) |
| 282 | +} |
278 | 283 |
|
279 | | - switch attrs := v.(type) { |
280 | | - case []slog.Attr: |
281 | | - *r = *r.WithContext(context.WithValue(r.Context(), customAttributesCtxKey, append(attrs, attr))) |
| 284 | +// AddContextAttributes is the same as AddCustomAttributes, but it doesn't need access to the request struct. |
| 285 | +func AddContextAttributes(ctx context.Context, attr slog.Attr) { |
| 286 | + if v := ctx.Value(customAttributesCtxKey); v != nil { |
| 287 | + if m, ok := v.(*sync.Map); ok { |
| 288 | + m.Store(attr.Key, attr.Value) |
| 289 | + } |
282 | 290 | } |
283 | 291 | } |
284 | 292 |
|
|
0 commit comments