Skip to content

Commit 0b9aac4

Browse files
haolingdong-msfthaolingdong-msft
authored
Add context in PollingStrategy to make OpenTelemetry tracing work (Azure#27926)
Add context in PollingStrategy to make OpenTelemetry tracing work
1 parent 736fb12 commit 0b9aac4

File tree

4 files changed

+56
-11
lines changed

4 files changed

+56
-11
lines changed

‎sdk/core/azure-core/src/main/java/com/azure/core/http/rest/RequestOptions.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ EnumSet<ErrorOptions> getErrorOptions() {
149149
*
150150
* @return The additional context that is passed during the service call.
151151
*/
152-
Context getContext() {
152+
public Context getContext() {
153153
return context;
154154
}
155155

‎sdk/core/azure-core/src/main/java/com/azure/core/util/polling/DefaultPollingStrategy.java‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import com.azure.core.http.HttpPipeline;
77
import com.azure.core.http.rest.Response;
88
import com.azure.core.implementation.serializer.DefaultJsonSerializer;
9+
import com.azure.core.util.Context;
910
import com.azure.core.util.serializer.JsonSerializer;
1011
import com.azure.core.util.serializer.TypeReference;
1112
import reactor.core.publisher.Mono;
@@ -33,7 +34,7 @@ public final class DefaultPollingStrategy<T, U> implements PollingStrategy<T, U>
3334
* @throws NullPointerException If {@code httpPipeline} is null.
3435
*/
3536
public DefaultPollingStrategy(HttpPipeline httpPipeline) {
36-
this(httpPipeline, new DefaultJsonSerializer());
37+
this(httpPipeline, new DefaultJsonSerializer(), Context.NONE);
3738
}
3839

3940
/**
@@ -46,9 +47,23 @@ public DefaultPollingStrategy(HttpPipeline httpPipeline) {
4647
* @throws NullPointerException If {@code httpPipeline} is null.
4748
*/
4849
public DefaultPollingStrategy(HttpPipeline httpPipeline, JsonSerializer serializer) {
50+
this(httpPipeline, serializer, Context.NONE);
51+
}
52+
53+
/**
54+
* Creates a chained polling strategy with 3 known polling strategies, {@link OperationResourcePollingStrategy},
55+
* {@link LocationPollingStrategy}, and {@link StatusCheckPollingStrategy}, in this order, with a custom
56+
* serializer.
57+
*
58+
* @param httpPipeline an instance of {@link HttpPipeline} to send requests with
59+
* @param serializer a custom serializer for serializing and deserializing polling responses
60+
* @param context an instance of {@link Context}
61+
* @throws NullPointerException If {@code httpPipeline} is null.
62+
*/
63+
public DefaultPollingStrategy(HttpPipeline httpPipeline, JsonSerializer serializer, Context context) {
4964
this.chainedPollingStrategy = new ChainedPollingStrategy<>(Arrays.asList(
50-
new OperationResourcePollingStrategy<>(httpPipeline, serializer, null),
51-
new LocationPollingStrategy<>(httpPipeline, serializer),
65+
new OperationResourcePollingStrategy<>(httpPipeline, serializer, null, context),
66+
new LocationPollingStrategy<>(httpPipeline, serializer, context),
5267
new StatusCheckPollingStrategy<>(serializer)));
5368
}
5469

‎sdk/core/azure-core/src/main/java/com/azure/core/util/polling/LocationPollingStrategy.java‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import com.azure.core.implementation.ImplUtils;
1414
import com.azure.core.implementation.serializer.DefaultJsonSerializer;
1515
import com.azure.core.util.BinaryData;
16+
import com.azure.core.util.Context;
1617
import com.azure.core.util.logging.ClientLogger;
1718
import com.azure.core.util.polling.implementation.PollingConstants;
1819
import com.azure.core.util.polling.implementation.PollingUtils;
@@ -40,6 +41,7 @@ public class LocationPollingStrategy<T, U> implements PollingStrategy<T, U> {
4041

4142
private final HttpPipeline httpPipeline;
4243
private final ObjectSerializer serializer;
44+
private final Context context;
4345

4446
/**
4547
* Creates an instance of the location polling strategy using a JSON serializer.
@@ -48,7 +50,7 @@ public class LocationPollingStrategy<T, U> implements PollingStrategy<T, U> {
4850
* @throws NullPointerException If {@code httpPipeline} is null.
4951
*/
5052
public LocationPollingStrategy(HttpPipeline httpPipeline) {
51-
this(httpPipeline, DEFAULT_SERIALIZER);
53+
this(httpPipeline, DEFAULT_SERIALIZER, Context.NONE);
5254
}
5355

5456
/**
@@ -59,8 +61,21 @@ public LocationPollingStrategy(HttpPipeline httpPipeline) {
5961
* @throws NullPointerException If {@code httpPipeline} is null.
6062
*/
6163
public LocationPollingStrategy(HttpPipeline httpPipeline, ObjectSerializer serializer) {
64+
this(httpPipeline, serializer, Context.NONE);
65+
}
66+
67+
/**
68+
* Creates an instance of the location polling strategy.
69+
*
70+
* @param httpPipeline an instance of {@link HttpPipeline} to send requests with
71+
* @param serializer a custom serializer for serializing and deserializing polling responses
72+
* @param context an instance of {@link Context}
73+
* @throws NullPointerException If {@code httpPipeline} is null.
74+
*/
75+
public LocationPollingStrategy(HttpPipeline httpPipeline, ObjectSerializer serializer, Context context) {
6276
this.httpPipeline = Objects.requireNonNull(httpPipeline, "'httpPipeline' cannot be null");
6377
this.serializer = (serializer == null) ? DEFAULT_SERIALIZER : serializer;
78+
this.context = context;
6479
}
6580

6681
@Override
@@ -108,7 +123,7 @@ public Mono<PollResponse<T>> onInitialResponse(Response<?> response, PollingCont
108123
@Override
109124
public Mono<PollResponse<T>> poll(PollingContext<T> pollingContext, TypeReference<T> pollResponseType) {
110125
HttpRequest request = new HttpRequest(HttpMethod.GET, pollingContext.getData(PollingConstants.LOCATION));
111-
return httpPipeline.send(request).flatMap(response -> {
126+
return httpPipeline.send(request, this.context).flatMap(response -> {
112127
HttpHeader locationHeader = response.getHeaders().get(PollingConstants.LOCATION);
113128
if (locationHeader != null) {
114129
pollingContext.setData(PollingConstants.LOCATION, locationHeader.getValue());
@@ -157,7 +172,7 @@ public Mono<U> getResult(PollingContext<T> pollingContext, TypeReference<U> resu
157172
return PollingUtils.deserializeResponse(BinaryData.fromString(latestResponseBody), serializer, resultType);
158173
} else {
159174
HttpRequest request = new HttpRequest(HttpMethod.GET, finalGetUrl);
160-
return httpPipeline.send(request)
175+
return httpPipeline.send(request, this.context)
161176
.flatMap(HttpResponse::getBodyAsByteArray)
162177
.map(BinaryData::fromBytes)
163178
.flatMap(binaryData -> PollingUtils.deserializeResponse(binaryData, serializer, resultType));

‎sdk/core/azure-core/src/main/java/com/azure/core/util/polling/OperationResourcePollingStrategy.java‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import com.azure.core.http.rest.Response;
1313
import com.azure.core.implementation.serializer.DefaultJsonSerializer;
1414
import com.azure.core.util.BinaryData;
15+
import com.azure.core.util.Context;
1516
import com.azure.core.util.polling.implementation.PollingConstants;
1617
import com.azure.core.util.polling.implementation.PollingUtils;
1718
import com.azure.core.util.serializer.ObjectSerializer;
@@ -37,6 +38,7 @@ public class OperationResourcePollingStrategy<T, U> implements PollingStrategy<T
3738
private final HttpPipeline httpPipeline;
3839
private final ObjectSerializer serializer;
3940
private final String operationLocationHeaderName;
41+
private final Context context;
4042

4143
/**
4244
* Creates an instance of the operation resource polling strategy using a JSON serializer and "Operation-Location"
@@ -45,24 +47,37 @@ public class OperationResourcePollingStrategy<T, U> implements PollingStrategy<T
4547
* @param httpPipeline an instance of {@link HttpPipeline} to send requests with
4648
*/
4749
public OperationResourcePollingStrategy(HttpPipeline httpPipeline) {
48-
this(httpPipeline, new DefaultJsonSerializer(), DEFAULT_OPERATION_LOCATION_HEADER);
50+
this(httpPipeline, new DefaultJsonSerializer(), DEFAULT_OPERATION_LOCATION_HEADER, Context.NONE);
4951
}
5052

5153
/**
5254
* Creates an instance of the operation resource polling strategy.
53-
*
5455
* @param httpPipeline an instance of {@link HttpPipeline} to send requests with
5556
* @param serializer a custom serializer for serializing and deserializing polling responses
5657
* @param operationLocationHeaderName a custom header for polling the long running operation
5758
*/
5859
public OperationResourcePollingStrategy(HttpPipeline httpPipeline, ObjectSerializer serializer,
5960
String operationLocationHeaderName) {
61+
this(httpPipeline, serializer, operationLocationHeaderName, Context.NONE);
62+
}
63+
64+
/**
65+
* Creates an instance of the operation resource polling strategy.
66+
* @param httpPipeline an instance of {@link HttpPipeline} to send requests with
67+
* @param serializer a custom serializer for serializing and deserializing polling responses
68+
* @param operationLocationHeaderName a custom header for polling the long running operation
69+
* @param context an instance of {@link com.azure.core.util.Context}
70+
*/
71+
public OperationResourcePollingStrategy(HttpPipeline httpPipeline, ObjectSerializer serializer,
72+
String operationLocationHeaderName, Context context) {
6073
this.httpPipeline = Objects.requireNonNull(httpPipeline, "'httpPipeline' cannot be null");
6174
this.serializer = serializer != null ? serializer : new DefaultJsonSerializer();
6275
this.operationLocationHeaderName = operationLocationHeaderName != null ? operationLocationHeaderName
6376
: DEFAULT_OPERATION_LOCATION_HEADER;
77+
this.context = context;
6478
}
6579

80+
6681
@Override
6782
public Mono<Boolean> canPoll(Response<?> initialResponse) {
6883
HttpHeader operationLocationHeader = initialResponse.getHeaders().get(operationLocationHeaderName);
@@ -111,7 +126,7 @@ public Mono<PollResponse<T>> onInitialResponse(Response<?> response, PollingCont
111126
@Override
112127
public Mono<PollResponse<T>> poll(PollingContext<T> pollingContext, TypeReference<T> pollResponseType) {
113128
HttpRequest request = new HttpRequest(HttpMethod.GET, pollingContext.getData(operationLocationHeaderName));
114-
return httpPipeline.send(request).flatMap(response -> response.getBodyAsByteArray()
129+
return httpPipeline.send(request, this.context).flatMap(response -> response.getBodyAsByteArray()
115130
.map(BinaryData::fromBytes)
116131
.flatMap(binaryData -> PollingUtils.deserializeResponse(
117132
binaryData, serializer, new TypeReference<PollResult>() { })
@@ -157,7 +172,7 @@ public Mono<U> getResult(PollingContext<T> pollingContext, TypeReference<U> resu
157172
return PollingUtils.deserializeResponse(BinaryData.fromString(latestResponseBody), serializer, resultType);
158173
} else {
159174
HttpRequest request = new HttpRequest(HttpMethod.GET, finalGetUrl);
160-
return httpPipeline.send(request)
175+
return httpPipeline.send(request, this.context)
161176
.flatMap(HttpResponse::getBodyAsByteArray)
162177
.map(BinaryData::fromBytes)
163178
.flatMap(binaryData -> PollingUtils.deserializeResponse(binaryData, serializer, resultType));

0 commit comments

Comments
 (0)