Skip to content

Commit eb24eb1

Browse files
elekelek
authored
http: support baggage for new/orphan traces (#61)
1 parent a1546ad commit eb24eb1

2 files changed

Lines changed: 277 additions & 13 deletions

File tree

‎http/info.go‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,19 @@ func TraceInfoFromHeader(header HeaderGetter, allowedBaggage ...string) (rv Trac
6868
traceParent := header.Get(traceParentHeader)
6969
traceState := header.Get(traceStateHeader)
7070
baggage := header.Get(baggageHeader)
71+
bm := map[string]string{}
72+
if baggage != "" {
73+
for _, kv := range strings.Split(baggage, ",") {
74+
if key, value, ok := strings.Cut(kv, "="); ok {
75+
for _, b := range allowedBaggage {
76+
if key == b {
77+
bm[key] = value
78+
break
79+
}
80+
}
81+
}
82+
}
83+
}
7184

7285
if traceParent != "" {
7386
parts := strings.Split(traceParent, "-")
@@ -91,19 +104,6 @@ func TraceInfoFromHeader(header HeaderGetter, allowedBaggage ...string) (rv Trac
91104
return rv
92105
}
93106

94-
bm := map[string]string{}
95-
if baggage != "" {
96-
for _, kv := range strings.Split(baggage, ",") {
97-
if key, value, ok := strings.Cut(kv, "="); ok {
98-
for _, b := range allowedBaggage {
99-
if key == b {
100-
bm[key] = value
101-
break
102-
}
103-
}
104-
}
105-
}
106-
}
107107
return TraceInfo{
108108
TraceId: &traceID,
109109
ParentId: &parentID,
@@ -116,6 +116,7 @@ func TraceInfoFromHeader(header HeaderGetter, allowedBaggage ...string) (rv Trac
116116
if strings.Contains(traceState, "sampled=true") {
117117
return TraceInfo{
118118
Sampled: true,
119+
Baggage: bm,
119120
}
120121
}
121122
return rv

‎http/server_test.go‎

Lines changed: 263 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,263 @@
1+
// Copyright (C) 2022 Storj Labs, Inc.
2+
// See LICENSE for copying information.
3+
4+
package http
5+
6+
import (
7+
"encoding/json"
8+
"fmt"
9+
"net/http"
10+
"net/http/httptest"
11+
"testing"
12+
13+
"github.com/spacemonkeygo/monkit/v3"
14+
)
15+
16+
type TraceResponse struct {
17+
TraceID string `json:"trace_id"`
18+
SpanID string `json:"span_id"`
19+
Annotations map[string]string `json:"annotations"`
20+
}
21+
22+
func TestTraceHandlerIntegration(t *testing.T) {
23+
scope := monkit.Package()
24+
25+
// Create a simple handler that returns trace information
26+
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
27+
span := monkit.SpanFromCtx(r.Context())
28+
if span == nil {
29+
http.Error(w, "No span found in context", http.StatusInternalServerError)
30+
return
31+
}
32+
33+
annotations := make(map[string]string)
34+
for _, annotation := range span.Annotations() {
35+
annotations[annotation.Name] = annotation.Value
36+
}
37+
38+
response := TraceResponse{
39+
TraceID: fmt.Sprintf("%016x", span.Trace().Id()),
40+
SpanID: fmt.Sprintf("%016x", span.Id()),
41+
Annotations: annotations,
42+
}
43+
44+
w.Header().Set("Content-Type", "application/json")
45+
json.NewEncoder(w).Encode(response)
46+
})
47+
48+
// Wrap with TraceHandler
49+
traceHandler := TraceHandler(handler, scope, "foo")
50+
51+
// Create test server
52+
server := httptest.NewServer(traceHandler)
53+
defer server.Close()
54+
55+
t.Run("propagate existing trace", func(t *testing.T) {
56+
req, err := http.NewRequest("GET", server.URL+"/test", nil)
57+
if err != nil {
58+
t.Fatalf("Failed to create request: %v", err)
59+
}
60+
61+
req.Header.Set("traceparent", "00-0000000000000001-00000002-01")
62+
req.Header.Set("baggage", "foo=bar,forbidden=ignore")
63+
64+
traceResp := doRequest(t, err, req)
65+
66+
if traceResp.TraceID == "" {
67+
t.Error("Expected trace ID to be present")
68+
}
69+
70+
if traceResp.SpanID == "" {
71+
t.Error("Expected span ID to be present")
72+
}
73+
74+
if traceResp.Annotations["http.uri"] != "/test" {
75+
t.Errorf("Expected http.uri annotation to be '/test', got '%s'", traceResp.Annotations["http.uri"])
76+
}
77+
78+
if traceResp.Annotations["foo"] != "bar" {
79+
t.Errorf("Annotation is missing")
80+
}
81+
82+
if traceResp.Annotations["forbidden"] != "" {
83+
t.Errorf("Annotation should be missing")
84+
}
85+
})
86+
87+
t.Run("orphan trace", func(t *testing.T) {
88+
req, err := http.NewRequest("GET", server.URL+"/test", nil)
89+
if err != nil {
90+
t.Fatalf("Failed to create request: %v", err)
91+
}
92+
93+
req.Header.Set("baggage", "foo=bar,forbidden=ignore")
94+
req.Header.Set("tracestate", "sampled=true")
95+
96+
traceResp := doRequest(t, err, req)
97+
98+
if traceResp.TraceID == "" {
99+
t.Error("Expected trace ID to be present")
100+
}
101+
102+
if traceResp.SpanID == "" {
103+
t.Error("Expected span ID to be present")
104+
}
105+
106+
if traceResp.Annotations["http.uri"] != "/test" {
107+
t.Errorf("Expected http.uri annotation to be '/test', got '%s'", traceResp.Annotations["http.uri"])
108+
}
109+
110+
if traceResp.Annotations["foo"] != "bar" {
111+
t.Errorf("Annotation is missing")
112+
}
113+
114+
if traceResp.Annotations["forbidden"] != "" {
115+
t.Errorf("Annotation should be missing")
116+
}
117+
})
118+
}
119+
120+
func doRequest(t *testing.T, err error, req *http.Request) TraceResponse {
121+
// Make request
122+
client := &http.Client{}
123+
resp, err := client.Do(req)
124+
if err != nil {
125+
t.Fatalf("Request failed: %v", err)
126+
}
127+
defer resp.Body.Close()
128+
129+
if resp.StatusCode != http.StatusOK {
130+
t.Fatalf("Expected status 200, got %d", resp.StatusCode)
131+
}
132+
133+
// Parse response
134+
var traceResp TraceResponse
135+
if err := json.NewDecoder(resp.Body).Decode(&traceResp); err != nil {
136+
t.Fatalf("Failed to decode response: %v", err)
137+
}
138+
return traceResp
139+
}
140+
141+
func TestTraceHandlerWithCustomBaggage(t *testing.T) {
142+
scope := monkit.Package()
143+
144+
// Create handler that checks for baggage annotations
145+
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
146+
span := monkit.SpanFromCtx(r.Context())
147+
if span == nil {
148+
http.Error(w, "No span found in context", http.StatusInternalServerError)
149+
return
150+
}
151+
152+
// The baggage should be added as annotations to the span
153+
// We'll just return success - the baggage testing is indirect through span annotations
154+
w.WriteHeader(http.StatusOK)
155+
w.Write([]byte("OK"))
156+
})
157+
158+
// Create TraceHandler with allowed baggage
159+
traceHandler := TraceHandler(handler, scope, "allowed-key", "another-allowed")
160+
161+
server := httptest.NewServer(traceHandler)
162+
defer server.Close()
163+
164+
// Test with baggage header
165+
req, err := http.NewRequest("GET", server.URL, nil)
166+
if err != nil {
167+
t.Fatalf("Failed to create request: %v", err)
168+
}
169+
170+
req.Header.Set("traceparent", "00-0000000000000001-00000002-01")
171+
req.Header.Set("baggage", "allowed-key=allowed-value,not-allowed=ignored,another-allowed=another-value")
172+
173+
client := &http.Client{}
174+
resp, err := client.Do(req)
175+
if err != nil {
176+
t.Fatalf("Request failed: %v", err)
177+
}
178+
defer resp.Body.Close()
179+
180+
if resp.StatusCode != http.StatusOK {
181+
t.Fatalf("Expected status 200, got %d", resp.StatusCode)
182+
}
183+
184+
// The test passes if the handler completes successfully
185+
// Baggage is verified indirectly through the span annotations in the TraceHandler
186+
}
187+
188+
func TestTraceHandlerContextPropagation(t *testing.T) {
189+
scope := monkit.Package()
190+
191+
// Handler that verifies context propagation
192+
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
193+
ctx := r.Context()
194+
span := monkit.SpanFromCtx(ctx)
195+
196+
if span == nil {
197+
http.Error(w, "No span in context", http.StatusInternalServerError)
198+
return
199+
}
200+
201+
// Try to start a child span to verify context works
202+
childCtx := ctx
203+
defer scope.Func().RemoteTrace(&childCtx, span.Id(), span.Trace())(nil)
204+
205+
childSpan := monkit.SpanFromCtx(childCtx)
206+
if childSpan == nil {
207+
http.Error(w, "Could not create child span", http.StatusInternalServerError)
208+
return
209+
}
210+
211+
response := map[string]string{
212+
"parent_trace_id": fmt.Sprintf("%016x", span.Trace().Id()),
213+
"parent_span_id": fmt.Sprintf("%016x", span.Id()),
214+
"child_trace_id": fmt.Sprintf("%016x", childSpan.Trace().Id()),
215+
"child_span_id": fmt.Sprintf("%016x", childSpan.Id()),
216+
}
217+
218+
w.Header().Set("Content-Type", "application/json")
219+
json.NewEncoder(w).Encode(response)
220+
})
221+
222+
traceHandler := TraceHandler(handler, scope)
223+
server := httptest.NewServer(traceHandler)
224+
defer server.Close()
225+
226+
// Make request with trace parent
227+
req, err := http.NewRequest("GET", server.URL, nil)
228+
if err != nil {
229+
t.Fatalf("Failed to create request: %v", err)
230+
}
231+
req.Header.Set("traceparent", "00-0000000000000001-00000002-01")
232+
233+
resp, err := http.DefaultClient.Do(req)
234+
if err != nil {
235+
t.Fatalf("Request failed: %v", err)
236+
}
237+
defer resp.Body.Close()
238+
239+
if resp.StatusCode != http.StatusOK {
240+
t.Fatalf("Expected status 200, got %d", resp.StatusCode)
241+
}
242+
243+
var result map[string]string
244+
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
245+
t.Fatalf("Failed to decode response: %v", err)
246+
}
247+
248+
// Verify trace IDs match
249+
if result["parent_trace_id"] != result["child_trace_id"] {
250+
t.Errorf("Parent and child should share same trace ID: parent=%s, child=%s",
251+
result["parent_trace_id"], result["child_trace_id"])
252+
}
253+
254+
// Verify span IDs are different
255+
if result["parent_span_id"] == result["child_span_id"] {
256+
t.Error("Parent and child should have different span IDs")
257+
}
258+
259+
// Verify expected trace ID
260+
if result["parent_trace_id"] != "0000000000000001" {
261+
t.Errorf("Expected trace ID 0000000000000001, got %s", result["parent_trace_id"])
262+
}
263+
}

0 commit comments

Comments
 (0)