Skip to content

Commit 49f3c4c

Browse files
jaydelucajaydeluca
committed
Add test and header fallback
1 parent 50c83ae commit 49f3c4c

File tree

3 files changed

+210
-1
lines changed

3 files changed

+210
-1
lines changed

‎instrumentation/spring/spring-webflux/spring-webflux-5.3/library/src/main/java/io/opentelemetry/instrumentation/spring/webflux/v5_3/WebfluxTextMapGetter.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ enum WebfluxTextMapGetter implements TextMapGetter<ServerWebExchange> {
1818

1919
@Override
2020
public Iterable<String> keys(ServerWebExchange exchange) {
21-
return exchange.getRequest().getHeaders().keySet();
21+
return HeaderUtil.getKeys(exchange.getRequest().getHeaders());
2222
}
2323

2424
@Nullable

‎instrumentation/spring/spring-webflux/spring-webflux-5.3/library/src/main/java/io/opentelemetry/instrumentation/spring/webflux/v5_3/internal/HeaderUtil.java‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,16 @@
66
package io.opentelemetry.instrumentation.spring.webflux.v5_3.internal;
77

88
import static java.util.Collections.emptyList;
9+
import static java.util.Collections.emptySet;
910

1011
import java.lang.invoke.MethodHandle;
1112
import java.lang.invoke.MethodHandles;
1213
import java.lang.invoke.MethodType;
14+
import java.util.HashSet;
1315
import java.util.List;
16+
import java.util.Map;
17+
import java.util.Set;
18+
import java.util.function.BiConsumer;
1419
import java.util.function.Supplier;
1520
import javax.annotation.Nullable;
1621
import org.springframework.http.HttpHeaders;
@@ -23,6 +28,8 @@
2328
public final class HeaderUtil {
2429

2530
@Nullable private static final MethodHandle GET_HEADERS;
31+
@Nullable private static final MethodHandle FOR_EACH;
32+
@Nullable private static final MethodHandle KEY_SET;
2633

2734
static {
2835
GET_HEADERS =
@@ -32,6 +39,33 @@ public final class HeaderUtil {
3239
() ->
3340
findGetHeadersMethod(
3441
MethodType.methodType(List.class, Object.class))); // before spring web 7.0
42+
43+
// Spring Web 7+
44+
MethodHandle forEach = null;
45+
try {
46+
forEach =
47+
MethodHandles.lookup()
48+
.findVirtual(
49+
HttpHeaders.class,
50+
"forEach",
51+
MethodType.methodType(void.class, BiConsumer.class));
52+
} catch (Throwable t) {
53+
// ignore - will fall back to keySet
54+
}
55+
FOR_EACH = forEach;
56+
57+
// Spring Web 6 and earlier
58+
MethodHandle keySet = null;
59+
if (FOR_EACH == null) {
60+
try {
61+
keySet =
62+
MethodHandles.lookup()
63+
.findVirtual(Map.class, "keySet", MethodType.methodType(Set.class));
64+
} catch (Throwable t) {
65+
// ignore
66+
}
67+
}
68+
KEY_SET = keySet;
3569
}
3670

3771
@Nullable
@@ -64,5 +98,27 @@ public static List<String> getHeader(HttpHeaders headers, String name) {
6498
return emptyList();
6599
}
66100

101+
@SuppressWarnings("unchecked") // HttpHeaders is a Map in Spring Web 6 and earlier
102+
public static Set<String> getKeys(HttpHeaders headers) {
103+
if (FOR_EACH != null) {
104+
// Spring Web 7: HttpHeaders has forEach(BiConsumer) method
105+
try {
106+
Set<String> keys = new HashSet<>();
107+
FOR_EACH.invoke(headers, (BiConsumer<String, ?>) (key, value) -> keys.add(key));
108+
return keys;
109+
} catch (Throwable t) {
110+
// ignore
111+
}
112+
} else if (KEY_SET != null) {
113+
// Spring Web 6 and earlier: HttpHeaders extends Map
114+
try {
115+
return (Set<String>) KEY_SET.invoke(headers);
116+
} catch (Throwable t) {
117+
// ignore
118+
}
119+
}
120+
return emptySet();
121+
}
122+
67123
private HeaderUtil() {}
68124
}

‎instrumentation/spring/spring-webflux/spring-webflux-5.3/library/src/test/java/io/opentelemetry/instrumentation/spring/webflux/v5_3/WebfluxTextMapGetterTest.java‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.spring.webflux.v5_3;
7+
8+
import static org.assertj.core.api.Assertions.assertThat;
9+
10+
import java.util.ArrayList;
11+
import java.util.List;
12+
import org.junit.jupiter.api.Test;
13+
import org.springframework.http.HttpHeaders;
14+
import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
15+
import org.springframework.mock.web.server.MockServerWebExchange;
16+
import org.springframework.web.server.ServerWebExchange;
17+
18+
class WebfluxTextMapGetterTest {
19+
20+
@Test
21+
void testKeysWithMultipleHeaders() {
22+
MockServerHttpRequest request =
23+
MockServerHttpRequest.get("/test")
24+
.header("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")
25+
.header("tracestate", "congo=t61rcWkgMzE")
26+
.header("custom-header", "custom-value")
27+
.header("x-forwarded-for", "192.168.1.1")
28+
.build();
29+
30+
MockServerWebExchange exchange = MockServerWebExchange.from(request);
31+
32+
Iterable<String> keys = WebfluxTextMapGetter.INSTANCE.keys(exchange);
33+
assertThat(keys).contains("traceparent", "tracestate", "custom-header", "x-forwarded-for");
34+
}
35+
36+
@Test
37+
void testGet() {
38+
MockServerHttpRequest request =
39+
MockServerHttpRequest.get("/test")
40+
.header("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")
41+
.header("custom-header", "custom-value")
42+
.build();
43+
44+
MockServerWebExchange exchange = MockServerWebExchange.from(request);
45+
46+
String traceparent = WebfluxTextMapGetter.INSTANCE.get(exchange, "traceparent");
47+
assertThat(traceparent).isEqualTo("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01");
48+
49+
String customHeader = WebfluxTextMapGetter.INSTANCE.get(exchange, "custom-header");
50+
assertThat(customHeader).isEqualTo("custom-value");
51+
}
52+
53+
@Test
54+
void testGetAll() {
55+
MockServerHttpRequest request =
56+
MockServerHttpRequest.get("/test")
57+
.header("accept", "application/json")
58+
.header("accept", "text/html")
59+
.build();
60+
61+
MockServerWebExchange exchange = MockServerWebExchange.from(request);
62+
63+
List<String> acceptHeaders = new ArrayList<>();
64+
WebfluxTextMapGetter.INSTANCE.getAll(exchange, "accept").forEachRemaining(acceptHeaders::add);
65+
66+
assertThat(acceptHeaders).containsExactly("application/json", "text/html");
67+
}
68+
69+
@Test
70+
void testGetAllSingleValue() {
71+
MockServerHttpRequest request =
72+
MockServerHttpRequest.get("/test").header("content-type", "application/json").build();
73+
74+
MockServerWebExchange exchange = MockServerWebExchange.from(request);
75+
76+
List<String> contentTypes = new ArrayList<>();
77+
WebfluxTextMapGetter.INSTANCE
78+
.getAll(exchange, "content-type")
79+
.forEachRemaining(contentTypes::add);
80+
81+
assertThat(contentTypes).containsExactly("application/json");
82+
}
83+
84+
@Test
85+
void testGetNullExchange() {
86+
String result = WebfluxTextMapGetter.INSTANCE.get(null, "any-header");
87+
assertThat(result).isNull();
88+
}
89+
90+
@Test
91+
void testKeysDirectlyOnHttpHeaders() {
92+
HttpHeaders headers = new HttpHeaders();
93+
headers.add("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01");
94+
headers.add("tracestate", "congo=t61rcWkgMzE");
95+
headers.add("custom-header", "custom-value");
96+
97+
MockServerHttpRequest request = MockServerHttpRequest.get("/test").headers(headers).build();
98+
ServerWebExchange exchange = MockServerWebExchange.from(request);
99+
100+
// The keys() method internally calls HttpHeaders.keySet()
101+
// This will throw NoSuchMethodError with Spring Web 7 if not properly handled
102+
Iterable<String> keys = WebfluxTextMapGetter.INSTANCE.keys(exchange);
103+
assertThat(keys).hasSize(3).contains("traceparent", "tracestate", "custom-header");
104+
}
105+
106+
@Test
107+
void testGetAllNullExchange() {
108+
assertThat(WebfluxTextMapGetter.INSTANCE.getAll(null, "any-header")).isExhausted();
109+
}
110+
111+
@Test
112+
void testKeysWithBaggageHeader() {
113+
// Test that baggage headers are properly returned by keys()
114+
MockServerHttpRequest request =
115+
MockServerHttpRequest.get("/test")
116+
.header("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")
117+
.header("baggage", "test-baggage-key-1=test-baggage-value-1")
118+
.build();
119+
120+
MockServerWebExchange exchange = MockServerWebExchange.from(request);
121+
122+
Iterable<String> keys = WebfluxTextMapGetter.INSTANCE.keys(exchange);
123+
assertThat(keys).contains("traceparent", "baggage");
124+
125+
String baggageValue = WebfluxTextMapGetter.INSTANCE.get(exchange, "baggage");
126+
assertThat(baggageValue).isEqualTo("test-baggage-key-1=test-baggage-value-1");
127+
}
128+
129+
@Test
130+
void testKeysWithMultipleBaggageHeaders() {
131+
// Test that multiple baggage headers are properly returned by keys()
132+
// The W3C Baggage propagator needs to iterate through all headers to find baggage entries
133+
MockServerHttpRequest request =
134+
MockServerHttpRequest.get("/test")
135+
.header("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")
136+
.header("baggage", "test-baggage-key-1=test-baggage-value-1")
137+
.header("baggage", "test-baggage-key-2=test-baggage-value-2")
138+
.header("x-custom", "custom-value")
139+
.build();
140+
141+
MockServerWebExchange exchange = MockServerWebExchange.from(request);
142+
143+
Iterable<String> keys = WebfluxTextMapGetter.INSTANCE.keys(exchange);
144+
assertThat(keys).contains("traceparent", "baggage", "x-custom");
145+
146+
List<String> baggageValues = new ArrayList<>();
147+
WebfluxTextMapGetter.INSTANCE.getAll(exchange, "baggage").forEachRemaining(baggageValues::add);
148+
149+
assertThat(baggageValues)
150+
.containsExactly(
151+
"test-baggage-key-1=test-baggage-value-1", "test-baggage-key-2=test-baggage-value-2");
152+
}
153+
}

0 commit comments

Comments
 (0)