Skip to content

Commit 7fceadc

Browse files
stickfigurestickfigure
committed
Allow us to manipulate otel spans
1 parent 5e13d9b commit 7fceadc

8 files changed

Lines changed: 54 additions & 20 deletions

File tree

‎src/main/java/com/googlecode/objectify/ObjectifyFactory.java‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import com.googlecode.objectify.impl.ObjectifyImpl;
2020
import com.googlecode.objectify.impl.ObjectifyOptions;
2121
import com.googlecode.objectify.impl.Registrar;
22+
import com.googlecode.objectify.impl.Spanipulator;
23+
import com.googlecode.objectify.impl.SpanipulatorImpl;
2224
import com.googlecode.objectify.impl.Transactor;
2325
import com.googlecode.objectify.impl.TypeUtils;
2426
import com.googlecode.objectify.impl.translate.Translators;
@@ -46,7 +48,7 @@
4648
import java.util.SortedSet;
4749
import java.util.TreeMap;
4850
import java.util.TreeSet;
49-
import java.util.function.Supplier;
51+
import java.util.function.Function;
5052
import java.util.stream.Collectors;
5153

5254
/**
@@ -558,13 +560,13 @@ public <T> Ref<T> ref(final T value) {
558560
/**
559561
* For internal use, hides the optionality of otel.
560562
*/
561-
public <T> T span(final String name, final Supplier<T> work) {
563+
public <T> T span(final String name, final Function<Spanipulator, T> work) {
562564
if (tracer == null) {
563-
return work.get();
565+
return work.apply(Spanipulator.NOOP);
564566
} else {
565567
final Span span = tracer.spanBuilder(name).setSpanKind(SpanKind.CLIENT).startSpan();
566568
try (final Scope scope = span.makeCurrent()) {
567-
return work.get();
569+
return work.apply(new SpanipulatorImpl(span));
568570
} finally {
569571
span.end();
570572
}

‎src/main/java/com/googlecode/objectify/impl/DeleterImpl.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ public Result<Void> keys(final Key<?>... keys) {
4040

4141
@Override
4242
public Result<Void> keys(final Iterable<? extends Key<?>> keys) {
43-
return ofy.factory().span("delete", () -> {
43+
return ofy.factory().span("delete", spanipulator -> {
4444
final List<com.google.cloud.datastore.Key> rawKeys = new ArrayList<>();
4545
for (final Key<?> key: keys)
4646
rawKeys.add(key.getRaw());
@@ -61,7 +61,7 @@ public Result<Void> entity(final Object entity) {
6161

6262
@Override
6363
public Result<Void> entities(final Iterable<?> entities) {
64-
return ofy.factory().span("delete", () -> {
64+
return ofy.factory().span("delete", spanipulator -> {
6565
final List<com.google.cloud.datastore.Key> keys = new ArrayList<>();
6666
for (final Object obj : entities)
6767
keys.add(ofy.factory().keys().anythingToRawKey(obj, ofy.getOptions().getNamespace()));

‎src/main/java/com/googlecode/objectify/impl/LoaderImpl.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ public <E> Map<Key<E>, E> refs(final Iterable<Ref<E>> refs) {
107107

108108
@Override
109109
public <E> LoadResult<E> key(final Key<E> key) {
110-
return ofy.factory().span("load", () -> {
110+
return ofy.factory().span("load", spanipulator -> {
111111
final LoadResult<E> result = new LoadResult<>(key, createLoadEngine().load(key));
112112
// The low level API is not async, so let's ensure work is finished in the span.
113113
result.now();
@@ -154,7 +154,7 @@ public <E> Map<Key<E>, E> values(final Object... values) {
154154

155155
@Override
156156
public <E> Map<Key<E>, E> values(final Iterable<?> values) {
157-
return ofy.factory().span("load", () -> {
157+
return ofy.factory().span("load", spanipulator -> {
158158

159159
// Do this in a separate pass so any errors converting keys will show up before we try loading something
160160
final List<Key<E>> keys = new ArrayList<>();

‎src/main/java/com/googlecode/objectify/impl/QueryEngine.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,8 @@
2626
*/
2727
@Slf4j
2828
@RequiredArgsConstructor
29-
public class QueryEngine
30-
{
31-
/** */
29+
public class QueryEngine {
30+
3231
private final LoaderImpl loader;
3332
private final AsyncDatastoreReaderWriter ds;
3433

‎src/main/java/com/googlecode/objectify/impl/QueryImpl.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -292,7 +292,7 @@ public String toString() {
292292
*/
293293
@Override
294294
public LoadResult<T> first() {
295-
return loader.ofy.factory().span("query", () -> {
295+
return loader.ofy.factory().span("query", spanipulator -> {
296296
// By the way, this is the same thing that PreparedQuery.asSingleEntity() does internally
297297
final Iterator<T> it = this.limit(1).iterator();
298298

@@ -307,14 +307,14 @@ public LoadResult<T> first() {
307307

308308
@Override
309309
public AggregationResult aggregate(final Aggregation... aggregations) {
310-
return loader.ofy.factory().span("aggregate", () -> {
310+
return loader.ofy.factory().span("aggregate", spanipulator -> {
311311
return loader.createQueryEngine().queryAggregations(this.actual.newKeyQuery(), aggregations);
312312
});
313313
}
314314

315315
@Override
316316
public AggregationResult aggregate(final AggregationBuilder<?>... aggregations) {
317-
return loader.ofy.factory().span("aggregate", () -> {
317+
return loader.ofy.factory().span("aggregate", spanipulator -> {
318318
return loader.createQueryEngine().queryAggregations(this.actual.newKeyQuery(), aggregations);
319319
});
320320
}
@@ -332,7 +332,7 @@ public QueryResultIterable<T> iterable() {
332332
*/
333333
@Override
334334
public QueryResults<T> iterator() {
335-
return loader.ofy.factory().span("query", () -> {
335+
return loader.ofy.factory().span("query", spanipulator -> {
336336
// This is a bit odd from a span perspective; how should we track the iteration, which happens outside the span?
337337

338338
if (!actual.getProjection().isEmpty())

‎src/main/java/com/googlecode/objectify/impl/SaverImpl.java‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,10 @@
1717
*
1818
* @author Jeff Schnitzer <jeff@infohazard.org>
1919
*/
20-
public class SaverImpl implements Saver
21-
{
22-
/** */
20+
public class SaverImpl implements Saver {
21+
2322
private final ObjectifyImpl ofy;
2423

25-
/** */
2624
public SaverImpl(final ObjectifyImpl ofy) {
2725
this.ofy = ofy;
2826
}
@@ -48,7 +46,7 @@ public <E> Result<Map<Key<E>, E>> entities(final E... entities) {
4846

4947
@Override
5048
public <E> Result<Map<Key<E>, E>> entities(final Iterable<E> entities) {
51-
return ofy.factory().span("save", () -> {
49+
return ofy.factory().span("save", spanipulator -> {
5250
final Result<Map<Key<E>, E>> result = ofy.createWriteEngine().save(entities);
5351
// The low level API is not async, so let's ensure work is finished in the span.
5452
result.now();

‎src/main/java/com/googlecode/objectify/impl/Spanipulator.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package com.googlecode.objectify.impl;
2+
3+
import io.opentelemetry.api.trace.Span;
4+
5+
import java.util.function.Consumer;
6+
7+
/**
8+
* Interface that allows span manipulation in optional otel tracing. If otel is not enabled,
9+
* pass the NOOP version and the manipulate callbacks will not be applied.
10+
*/
11+
public interface Spanipulator {
12+
Spanipulator NOOP = new Spanipulator() {
13+
@Override
14+
public void update(final Consumer<Span> manipulate) {}
15+
};
16+
17+
/** Call this to update a span */
18+
void update(Consumer<Span> manipulate);
19+
}

‎src/main/java/com/googlecode/objectify/impl/SpanipulatorImpl.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package com.googlecode.objectify.impl;
2+
3+
import io.opentelemetry.api.trace.Span;
4+
import lombok.RequiredArgsConstructor;
5+
6+
import java.util.function.Consumer;
7+
8+
@RequiredArgsConstructor
9+
public class SpanipulatorImpl implements Spanipulator {
10+
private final Span span;
11+
12+
@Override
13+
public void update(final Consumer<Span> manipulate) {
14+
manipulate.accept(span);
15+
}
16+
}

0 commit comments

Comments
 (0)