27
27
import io .grpc .ServerCallHandler ;
28
28
import io .grpc .ServerInterceptor ;
29
29
import io .grpc .Status ;
30
+ import io .opentelemetry .api .GlobalOpenTelemetry ;
31
+ import io .opentelemetry .api .common .AttributeKey ;
32
+ import io .opentelemetry .api .trace .propagation .W3CTraceContextPropagator ;
33
+ import io .opentelemetry .context .propagation .ContextPropagators ;
34
+ import io .opentelemetry .sdk .OpenTelemetrySdk ;
35
+ import io .opentelemetry .sdk .testing .exporter .InMemorySpanExporter ;
36
+ import io .opentelemetry .sdk .trace .SdkTracerProvider ;
37
+ import io .opentelemetry .sdk .trace .data .SpanData ;
38
+ import io .opentelemetry .sdk .trace .export .SimpleSpanProcessor ;
30
39
import java .util .ArrayList ;
31
40
import java .util .List ;
32
41
import java .util .Map ;
35
44
import java .util .concurrent .ConcurrentHashMap ;
36
45
import java .util .concurrent .CopyOnWriteArrayList ;
37
46
import java .util .regex .Matcher ;
47
+ import org .junit .AfterClass ;
48
+ import org .junit .BeforeClass ;
38
49
import org .junit .Test ;
39
50
import org .junit .runner .RunWith ;
40
51
import org .junit .runners .JUnit4 ;
@@ -68,6 +79,13 @@ public static class ServerHeaderEnforcer implements ServerInterceptor {
68
79
private Map <String , CopyOnWriteArrayList <XGoogSpannerRequestId >> streamingResults ;
69
80
private List <String > gotValues ;
70
81
private Set <String > checkMethods ;
82
+ private static InMemorySpanExporter spanExporter ;
83
+ private static OpenTelemetrySdk openTelemetry ;
84
+
85
+ // This variable when set to true will allow checking for span
86
+ // attributes like `x_goog_spanner_request_id`, of which that
87
+ // requires considerable plumbing in this package.
88
+ private static final boolean ENABLED_USER_DEFINED_TRACES = false ;
71
89
72
90
ServerHeaderEnforcer (Set <String > checkMethods ) {
73
91
this .gotValues = new CopyOnWriteArrayList <String >();
@@ -76,6 +94,38 @@ public static class ServerHeaderEnforcer implements ServerInterceptor {
76
94
this .streamingResults =
77
95
new ConcurrentHashMap <String , CopyOnWriteArrayList <XGoogSpannerRequestId >>();
78
96
this .checkMethods = checkMethods ;
97
+ ServerHeaderEnforcer .setupOpenTelemetry ();
98
+ }
99
+
100
+ @ BeforeClass
101
+ public static void setupOpenTelemetry () {
102
+ SpannerOptions .resetActiveTracingFramework ();
103
+ SpannerOptions .enableOpenTelemetryTraces ();
104
+ GlobalOpenTelemetry .resetForTest ();
105
+
106
+ spanExporter = InMemorySpanExporter .create ();
107
+
108
+ SdkTracerProvider tracerProvider =
109
+ SdkTracerProvider .builder ()
110
+ .addSpanProcessor (SimpleSpanProcessor .create (spanExporter ))
111
+ .build ();
112
+
113
+ openTelemetry =
114
+ OpenTelemetrySdk .builder ()
115
+ .setPropagators (ContextPropagators .create (W3CTraceContextPropagator .getInstance ()))
116
+ .setTracerProvider (tracerProvider )
117
+ .buildAndRegisterGlobal ();
118
+ }
119
+
120
+ @ AfterClass
121
+ public static void closeOpenTelemetry () {
122
+ if (openTelemetry != null ) {
123
+ openTelemetry .close ();
124
+ }
125
+ }
126
+
127
+ public List <SpanData > getFinishedSpans () {
128
+ return this .spanExporter .getFinishedSpanItems ();
79
129
}
80
130
81
131
@ Override
@@ -108,6 +158,15 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
108
158
// Firstly assert and validate that at least we've got a requestId.
109
159
Matcher m = XGoogSpannerRequestId .REGEX .matcher (gotReqIdStr );
110
160
assertTrue (m .matches ());
161
+ if (ENABLED_USER_DEFINED_TRACES ) {
162
+ List <SpanData > spans = this .getFinishedSpans ();
163
+ assertNotNull (spans );
164
+ SpanData span = spans .get (0 );
165
+ String gotReqIdFromSpanAttribute =
166
+ span .getAttributes ().get (AttributeKey .stringKey ("x_goog_spanner_request_id" ));
167
+ assertNotNull (gotReqIdFromSpanAttribute );
168
+ assertEquals (gotReqIdStr , gotReqIdFromSpanAttribute );
169
+ }
111
170
112
171
XGoogSpannerRequestId reqId = XGoogSpannerRequestId .of (gotReqIdStr );
113
172
if (!saver .containsKey (methodName )) {
0 commit comments