|
20 | 20 | import datadog.trace.api.gateway.Flow; |
21 | 21 | import datadog.trace.api.gateway.Flow.Action.RequestBlockingAction; |
22 | 22 | import datadog.trace.api.gateway.IGSpanInfo; |
| 23 | +import datadog.trace.api.gateway.InferredProxySpan; |
23 | 24 | import datadog.trace.api.gateway.RequestContext; |
24 | 25 | import datadog.trace.api.gateway.RequestContextSlot; |
25 | 26 | import datadog.trace.api.naming.SpanNaming; |
@@ -147,21 +148,32 @@ public Context startSpan(REQUEST_CARRIER carrier, Context parentContext) { |
147 | 148 | instrumentationNames != null && instrumentationNames.length > 0 |
148 | 149 | ? instrumentationNames[0] |
149 | 150 | : DEFAULT_INSTRUMENTATION_NAME; |
150 | | - AgentSpanContext.Extracted extracted = |
151 | | - callIGCallbackStart(getExtractedSpanContext(parentContext)); |
| 151 | + AgentSpanContext extracted = getExtractedSpanContext(parentContext); |
| 152 | + // Call IG callbacks |
| 153 | + extracted = callIGCallbackStart(extracted); |
| 154 | + // Create gateway inferred span if needed |
| 155 | + extracted = startInferredProxySpan(parentContext, extracted); |
152 | 156 | AgentSpan span = |
153 | 157 | tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true); |
| 158 | + // Apply RequestBlockingAction if any |
154 | 159 | Flow<Void> flow = callIGCallbackRequestHeaders(span, carrier); |
155 | 160 | if (flow.getAction() instanceof RequestBlockingAction) { |
156 | 161 | span.setRequestBlockingAction((RequestBlockingAction) flow.getAction()); |
157 | 162 | } |
158 | | - AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter(); |
159 | | - if (null != carrier && null != getter) { |
160 | | - tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer()); |
161 | | - } |
| 163 | + // DSM Checkpoint |
| 164 | + tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer()); |
162 | 165 | return parentContext.with(span); |
163 | 166 | } |
164 | 167 |
|
| 168 | + protected AgentSpanContext startInferredProxySpan(Context context, AgentSpanContext extracted) { |
| 169 | + InferredProxySpan span; |
| 170 | + if (!Config.get().isInferredProxyPropagationEnabled() |
| 171 | + || (span = InferredProxySpan.fromContext(context)) == null) { |
| 172 | + return extracted; |
| 173 | + } |
| 174 | + return span.start(extracted); |
| 175 | + } |
| 176 | + |
165 | 177 | public AgentSpan onRequest( |
166 | 178 | final AgentSpan span, |
167 | 179 | final CONNECTION connection, |
@@ -390,8 +402,7 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) { |
390 | 402 | return span; |
391 | 403 | } |
392 | 404 |
|
393 | | - private AgentSpanContext.Extracted callIGCallbackStart( |
394 | | - @Nullable final AgentSpanContext.Extracted extracted) { |
| 405 | + private AgentSpanContext callIGCallbackStart(@Nullable final AgentSpanContext extracted) { |
395 | 406 | AgentTracer.TracerAPI tracer = tracer(); |
396 | 407 | Supplier<Flow<Object>> startedCbAppSec = |
397 | 408 | tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted()); |
@@ -527,10 +538,20 @@ private Flow<Void> callIGCallbackURI( |
527 | 538 |
|
528 | 539 | @Override |
529 | 540 | public AgentSpan beforeFinish(AgentSpan span) { |
| 541 | + // TODO Migrate beforeFinish to Context API |
530 | 542 | onRequestEndForInstrumentationGateway(span); |
| 543 | + // Close Serverless Gateway Inferred Span if any |
| 544 | + // finishInferredProxySpan(context); |
531 | 545 | return super.beforeFinish(span); |
532 | 546 | } |
533 | 547 |
|
| 548 | + protected void finishInferredProxySpan(Context context) { |
| 549 | + InferredProxySpan span; |
| 550 | + if ((span = InferredProxySpan.fromContext(context)) != null) { |
| 551 | + span.finish(); |
| 552 | + } |
| 553 | + } |
| 554 | + |
534 | 555 | private void onRequestEndForInstrumentationGateway(@Nonnull final AgentSpan span) { |
535 | 556 | if (span.getLocalRootSpan() != span) { |
536 | 557 | return; |
|
0 commit comments