Skip to content

Commit d02d959

Browse files
jimit-j-shahjimit-j-shah
committed
fix: Fixed Span nesting for ReadWriteTransactionCallable by using parent SpanContext instead of just parent Context (#1495)
* fix: Fixed the TraceUtil.startSpan method to use `SpanContext` for linking with the parent instead of `Context`. - This fixes the hierarchy of Spans appearing in a transaction under a Run method. - Tested using existing transaction test * Fixed commit reordering and typos * fix: lint errors * fix: Refactored the ReadWriteTransactioncallable.call method to use startSpan idiomatically - TraceUtil.startSpan needs more debugging - return DefaultTracerProvider instance (no-op) when initializing DisabledTraceUtil - this fixes the unit tests in DatastoreTest.testRunInTransactionWithReadWriteOption * feat: Added tracing for Transaction.RunQuery (#1499) * feat: Added span for Transactional RunQuery - tested * fix: lint * fix: patch apply issues * fix: refactor using boolean flag * fix: s/startSpan/startSpanWithParentContext
1 parent 6af5223 commit d02d959

File tree

7 files changed

+193
-46
lines changed

7 files changed

+193
-46
lines changed

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java

Lines changed: 61 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,9 @@
2525
import com.google.cloud.ServiceOptions;
2626
import com.google.cloud.datastore.execution.AggregationQueryExecutor;
2727
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
28-
import com.google.cloud.datastore.telemetry.TraceUtil.Context;
2928
import com.google.common.base.MoreObjects;
3029
import com.google.common.base.Preconditions;
30+
import com.google.common.base.Throwables;
3131
import com.google.common.collect.AbstractIterator;
3232
import com.google.common.collect.ImmutableList;
3333
import com.google.common.collect.ImmutableMap;
@@ -42,6 +42,11 @@
4242
import io.opencensus.common.Scope;
4343
import io.opencensus.trace.Span;
4444
import io.opencensus.trace.Status;
45+
import io.opentelemetry.api.common.Attributes;
46+
import io.opentelemetry.api.trace.SpanBuilder;
47+
import io.opentelemetry.api.trace.SpanKind;
48+
import io.opentelemetry.api.trace.StatusCode;
49+
import io.opentelemetry.context.Context;
4550
import java.util.ArrayList;
4651
import java.util.Arrays;
4752
import java.util.Collections;
@@ -53,7 +58,7 @@
5358
import java.util.Optional;
5459
import java.util.Set;
5560
import java.util.concurrent.Callable;
56-
import javax.annotation.Nonnull;
61+
import javax.annotation.Nullable;
5762

5863
final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datastore {
5964

@@ -106,15 +111,18 @@ static class ReadWriteTransactionCallable<T> implements Callable<T> {
106111
private volatile TransactionOptions options;
107112
private volatile Transaction transaction;
108113

114+
private final com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext;
115+
109116
ReadWriteTransactionCallable(
110117
Datastore datastore,
111118
TransactionCallable<T> callable,
112119
TransactionOptions options,
113-
@Nonnull Context parentTraceContext) {
120+
@Nullable com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext) {
114121
this.datastore = datastore;
115122
this.callable = callable;
116123
this.options = options;
117124
this.transaction = null;
125+
this.parentSpanContext = parentSpanContext;
118126
}
119127

120128
Datastore getDatastore() {
@@ -135,26 +143,57 @@ void setPrevTransactionId(ByteString transactionId) {
135143
options = options.toBuilder().setReadWrite(readWrite).build();
136144
}
137145

146+
private io.opentelemetry.api.trace.Span startSpanWithParentContext(
147+
String spanName,
148+
com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext) {
149+
com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil =
150+
datastore.getOptions().getTraceUtil();
151+
SpanBuilder spanBuilder =
152+
otelTraceUtil
153+
.getTracer()
154+
.spanBuilder(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN)
155+
.setSpanKind(SpanKind.PRODUCER)
156+
.setParent(
157+
Context.current()
158+
.with(
159+
io.opentelemetry.api.trace.Span.wrap(
160+
parentSpanContext.getSpanContext())));
161+
return spanBuilder.startSpan();
162+
}
163+
138164
@Override
139165
public T call() throws DatastoreException {
140-
com.google.cloud.datastore.telemetry.TraceUtil traceUtil =
141-
datastore.getOptions().getTraceUtil();
142-
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
143-
traceUtil.startSpan(
166+
// TODO Instead of using OTel Spans directly, TraceUtil.Span should be used here. However,
167+
// the same code in startSpanInternal doesn't work when EnabledTraceUtil.StartSpan is called
168+
// probably because of some thread-local caching that is getting lost. This needs more
169+
// debugging. The code below works and is idiomatic but could be prettier and more consistent
170+
// with the use of TraceUtil-provided framework.
171+
io.opentelemetry.api.trace.Span span =
172+
startSpanWithParentContext(
144173
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN,
145-
datastore.getOptions().getTraceUtil().getCurrentContext());
146-
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
174+
parentSpanContext);
175+
try (io.opentelemetry.context.Scope ignored = span.makeCurrent()) {
147176
transaction = datastore.newTransaction(options);
148177
T value = callable.run(transaction);
149178
transaction.commit();
150179
return value;
151180
} catch (Exception ex) {
152181
transaction.rollback();
182+
span.setStatus(StatusCode.ERROR, ex.getMessage());
183+
span.recordException(
184+
ex,
185+
Attributes.builder()
186+
.put("exception.message", ex.getMessage())
187+
.put("exception.type", ex.getClass().getName())
188+
.put("exception.stacktrace", Throwables.getStackTraceAsString(ex))
189+
.build());
190+
span.end();
153191
throw DatastoreException.propagateUserException(ex);
154192
} finally {
155193
if (transaction.isActive()) {
156194
transaction.rollback();
157195
}
196+
span.end();
158197
if (options != null
159198
&& options.getModeCase().equals(TransactionOptions.ModeCase.READ_WRITE)) {
160199
setPrevTransactionId(transaction.getTransactionId());
@@ -165,42 +204,30 @@ public T call() throws DatastoreException {
165204

166205
@Override
167206
public <T> T runInTransaction(final TransactionCallable<T> callable) {
168-
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
169-
otelTraceUtil.startSpan(
170-
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN);
171-
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
207+
try {
172208
return RetryHelper.runWithRetries(
173209
new ReadWriteTransactionCallable<T>(
174-
this, callable, null, otelTraceUtil.getCurrentContext()),
210+
this, callable, null, otelTraceUtil.getCurrentSpanContext()),
175211
retrySettings,
176212
TRANSACTION_EXCEPTION_HANDLER,
177213
getOptions().getClock());
178214
} catch (RetryHelperException e) {
179-
span.end(e);
180215
throw DatastoreException.translateAndThrow(e);
181-
} finally {
182-
span.end();
183216
}
184217
}
185218

186219
@Override
187220
public <T> T runInTransaction(
188221
final TransactionCallable<T> callable, TransactionOptions transactionOptions) {
189-
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
190-
otelTraceUtil.startSpan(
191-
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN);
192-
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
222+
try {
193223
return RetryHelper.runWithRetries(
194224
new ReadWriteTransactionCallable<T>(
195-
this, callable, transactionOptions, otelTraceUtil.getCurrentContext()),
225+
this, callable, transactionOptions, otelTraceUtil.getCurrentSpanContext()),
196226
retrySettings,
197227
TRANSACTION_EXCEPTION_HANDLER,
198228
getOptions().getClock());
199229
} catch (RetryHelperException e) {
200-
span.end(e);
201230
throw DatastoreException.translateAndThrow(e);
202-
} finally {
203-
span.end();
204231
}
205232
}
206233

@@ -258,11 +285,14 @@ public AggregationResults runAggregation(
258285

259286
com.google.datastore.v1.RunQueryResponse runQuery(
260287
final com.google.datastore.v1.RunQueryRequest requestPb) {
261-
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
262-
otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY);
263288
ReadOptions readOptions = requestPb.getReadOptions();
264-
span.setAttribute(
265-
"isTransactional", readOptions.hasTransaction() || readOptions.hasNewTransaction());
289+
boolean isTransactional = readOptions.hasTransaction() || readOptions.hasNewTransaction();
290+
String spanName =
291+
(isTransactional
292+
? com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY
293+
: com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY);
294+
com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName);
295+
span.setAttribute("isTransactional", isTransactional);
266296
span.setAttribute("readConsistency", readOptions.getReadConsistency().toString());
267297

268298
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
@@ -275,7 +305,7 @@ com.google.datastore.v1.RunQueryResponse runQuery(
275305
: TRANSACTION_OPERATION_EXCEPTION_HANDLER,
276306
getOptions().getClock());
277307
span.addEvent(
278-
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY + ": Completed",
308+
spanName + ": Completed",
279309
new ImmutableMap.Builder<String, Object>()
280310
.put("Received", response.getBatch().getEntityResultsCount())
281311
.put("More results", response.getBatch().getMoreResults().toString())
@@ -689,7 +719,7 @@ com.google.datastore.v1.BeginTransactionResponse beginTransaction(
689719
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
690720
otelTraceUtil.startSpan(
691721
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_BEGIN_TRANSACTION,
692-
otelTraceUtil.getCurrentContext());
722+
otelTraceUtil.getCurrentSpanContext());
693723
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope scope = span.makeCurrent()) {
694724
return RetryHelper.runWithRetries(
695725
() -> datastoreRpc.beginTransaction(requestPb),

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,11 @@
1919
import com.google.api.core.ApiFunction;
2020
import com.google.api.core.ApiFuture;
2121
import com.google.api.core.InternalApi;
22+
import com.google.cloud.datastore.telemetry.TraceUtil.SpanContext;
2223
import io.grpc.ManagedChannelBuilder;
24+
import io.opentelemetry.api.trace.Span;
25+
import io.opentelemetry.api.trace.Tracer;
26+
import io.opentelemetry.api.trace.TracerProvider;
2327
import java.util.Map;
2428
import javax.annotation.Nonnull;
2529
import javax.annotation.Nullable;
@@ -31,6 +35,13 @@
3135
@InternalApi
3236
public class DisabledTraceUtil implements TraceUtil {
3337

38+
static class SpanContext implements TraceUtil.SpanContext {
39+
@Override
40+
public io.opentelemetry.api.trace.SpanContext getSpanContext() {
41+
return null;
42+
}
43+
}
44+
3445
static class Span implements TraceUtil.Span {
3546
@Override
3647
public void end() {}
@@ -66,6 +77,10 @@ public TraceUtil.Span setAttribute(String key, boolean value) {
6677
return this;
6778
}
6879

80+
public io.opentelemetry.api.trace.Span getSpan() {
81+
return null;
82+
}
83+
6984
@Override
7085
public Scope makeCurrent() {
7186
return new Scope();
@@ -96,7 +111,7 @@ public Span startSpan(String spanName) {
96111
}
97112

98113
@Override
99-
public TraceUtil.Span startSpan(String spanName, TraceUtil.Context parent) {
114+
public TraceUtil.Span startSpan(String spanName, TraceUtil.SpanContext parentSpanContext) {
100115
return new Span();
101116
}
102117

@@ -111,4 +126,15 @@ public TraceUtil.Span getCurrentSpan() {
111126
public TraceUtil.Context getCurrentContext() {
112127
return new Context();
113128
}
129+
130+
@Nonnull
131+
@Override
132+
public TraceUtil.SpanContext getCurrentSpanContext() {
133+
return new SpanContext();
134+
}
135+
136+
@Override
137+
public Tracer getTracer() {
138+
return TracerProvider.noop().get(LIBRARY_NAME);
139+
}
114140
}

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,19 @@
2222
import com.google.api.core.ApiFutures;
2323
import com.google.api.core.InternalApi;
2424
import com.google.cloud.datastore.DatastoreOptions;
25+
import com.google.cloud.datastore.telemetry.TraceUtil.SpanContext;
2526
import com.google.common.base.Throwables;
2627
import io.grpc.ManagedChannelBuilder;
2728
import io.opentelemetry.api.GlobalOpenTelemetry;
2829
import io.opentelemetry.api.OpenTelemetry;
2930
import io.opentelemetry.api.common.Attributes;
3031
import io.opentelemetry.api.common.AttributesBuilder;
32+
import io.opentelemetry.api.trace.Span;
3133
import io.opentelemetry.api.trace.SpanBuilder;
3234
import io.opentelemetry.api.trace.SpanKind;
3335
import io.opentelemetry.api.trace.StatusCode;
3436
import io.opentelemetry.api.trace.Tracer;
37+
import io.opentelemetry.context.Context;
3538
import java.util.Map;
3639
import javax.annotation.Nonnull;
3740
import javax.annotation.Nullable;
@@ -70,6 +73,19 @@ public ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> getChannelConfi
7073
return null;
7174
}
7275

76+
static class SpanContext implements TraceUtil.SpanContext {
77+
private final io.opentelemetry.api.trace.SpanContext spanContext;
78+
79+
public SpanContext(io.opentelemetry.api.trace.SpanContext spanContext) {
80+
this.spanContext = spanContext;
81+
}
82+
83+
@Override
84+
public io.opentelemetry.api.trace.SpanContext getSpanContext() {
85+
return this.spanContext;
86+
}
87+
}
88+
7389
static class Span implements TraceUtil.Span {
7490
private final io.opentelemetry.api.trace.Span span;
7591
private final String spanName;
@@ -181,6 +197,10 @@ public TraceUtil.Span setAttribute(String key, boolean value) {
181197
return this;
182198
}
183199

200+
public io.opentelemetry.api.trace.Span getSpan() {
201+
return this.span;
202+
}
203+
184204
@Override
185205
public Scope makeCurrent() {
186206
try (io.opentelemetry.context.Scope scope = span.makeCurrent()) {
@@ -286,13 +306,15 @@ public Span startSpan(String spanName) {
286306
}
287307

288308
@Override
289-
public TraceUtil.Span startSpan(String spanName, TraceUtil.Context parent) {
290-
assert (parent instanceof Context);
309+
public TraceUtil.Span startSpan(String spanName, TraceUtil.SpanContext parentSpanContext) {
291310
SpanBuilder spanBuilder =
292311
tracer
293312
.spanBuilder(spanName)
294313
.setSpanKind(SpanKind.PRODUCER)
295-
.setParent(((Context) parent).context);
314+
.setParent(
315+
io.opentelemetry.context.Context.current()
316+
.with(
317+
io.opentelemetry.api.trace.Span.wrap(parentSpanContext.getSpanContext())));
296318
io.opentelemetry.api.trace.Span span =
297319
addSettingsAttributesToCurrentSpan(spanBuilder).startSpan();
298320
return new Span(span, spanName);
@@ -309,4 +331,15 @@ public TraceUtil.Span getCurrentSpan() {
309331
public TraceUtil.Context getCurrentContext() {
310332
return new Context(io.opentelemetry.context.Context.current());
311333
}
334+
335+
@Nonnull
336+
@Override
337+
public TraceUtil.SpanContext getCurrentSpanContext() {
338+
return new SpanContext(io.opentelemetry.api.trace.Span.current().getSpanContext());
339+
}
340+
341+
@Override
342+
public Tracer getTracer() {
343+
return this.tracer;
344+
}
312345
}

0 commit comments

Comments
 (0)