FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

fix: update otel integration to properly activate span context for lazy RPCs such as reads & writes by BenWhitehead · Pull Request #3255 · googleapis/java-storage · GitHub

Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension .java  (1) All 1 file type selected
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -64,11 +64,13 @@
import org.checkerframework.checker.nullness.qual.Nullable;

@SuppressWarnings("DuplicatedCode")
final class OtelStorageDecorator implements Storage {
public final class OtelStorageDecorator implements Storage {

/** Becomes the {@code otel.scope.name} attribute in a span */
private static final String OTEL_SCOPE_NAME = "cloud.google.com/java/storage";

private static final String BLOB_READ_SESSION = "blobReadSession";

@VisibleForTesting final Storage delegate;
private final OpenTelemetry otel;
private final Attributes baseAttributes;
Expand Down Expand Up @@ -1434,13 +1436,11 @@ public BlobWriteSession blobWriteSession(BlobInfo blobInfo, BlobWriteOption... o
.startSpan();
try (Scope ignore = sessionSpan.makeCurrent()) {
BlobWriteSession session = delegate.blobWriteSession(blobInfo, options);
return new OtelDecoratedBlobWriteSession(session);
return new OtelDecoratedBlobWriteSession(session, sessionSpan);
} catch (Throwable t) {
sessionSpan.recordException(t);
sessionSpan.setStatus(StatusCode.ERROR, t.getClass().getSimpleName());
throw t;
} finally {
sessionSpan.end();
}
}

Expand All @@ -1467,12 +1467,12 @@ public Blob moveBlob(MoveBlobRequest request) {
public ApiFuture<BlobReadSession> blobReadSession(BlobId id, BlobSourceOption... options) {
Span blobReadSessionSpan =
tracer
.spanBuilder("blobReadSession")
.spanBuilder(BLOB_READ_SESSION)
.setAttribute("gsutil.uri", id.toGsUtilUriWithGeneration())
.startSpan();
try (Scope ignore1 = blobReadSessionSpan.makeCurrent()) {
Context blobReadSessionContext = Context.current();
Span ready = tracer.spanBuilder("blobReadSession/ready").startSpan();
Span ready = tracer.spanBuilder(BLOB_READ_SESSION + "/ready").startSpan();
ApiFuture<BlobReadSession> blobReadSessionApiFuture = delegate.blobReadSession(id, options);
ApiFuture<BlobReadSession> futureDecorated =
ApiFutures.transform(
Expand Down Expand Up @@ -1561,7 +1561,7 @@ static UnaryOperator<RetryContext> retryContextDecorator(OpenTelemetry otel) {
return String.format(Locale.US, "gs://%s/", bucket);
}

private static final class TracerDecorator implements Tracer {
public static final class TracerDecorator implements Tracer {
@Nullable private final Context parentContextOverride;
private final Tracer delegate;
private final Attributes baseAttributes;
Expand All @@ -1578,7 +1578,7 @@ private TracerDecorator(
this.spanNamePrefix = spanNamePrefix;
}

private static TracerDecorator decorate(
public static TracerDecorator decorate(
@Nullable Context parentContextOverride,
OpenTelemetry otel,
Attributes baseAttributes,
Expand Down Expand Up @@ -1608,13 +1608,16 @@ static final class OtelDecoratedReadChannel implements ReadChannel {
@VisibleForTesting final ReadChannel reader;
private final Span span;

private volatile Scope scope;

private OtelDecoratedReadChannel(ReadChannel reader, Span span) {
this.reader = reader;
this.span = span;
}

@Override
public void seek(long position) throws IOException {
clearScope();
reader.seek(position);
}

Expand All @@ -1630,6 +1633,7 @@ public RestorableState<ReadChannel> capture() {

@Override
public ReadChannel limit(long limit) {
clearScope();
return reader.limit(limit);
}

Expand All @@ -1640,6 +1644,7 @@ public long limit() {

@Override
public int read(ByteBuffer dst) throws IOException {
setScope();
return reader.read(dst);
}

Expand All @@ -1650,21 +1655,38 @@ public boolean isOpen() {

@Override
public void close() {
setScope();
try {
reader.close();
} finally {
span.end();
clearScope();
}
}

private void clearScope() {
try (Scope ignore = scope) {
scope = null;
}
}

public void setScope() {
if (scope != null) {
clearScope();
}
scope = span.makeCurrent();
}
Comment on lines +1673 to 1678

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

This looks great! TIL to have explicit scope handling added to properly activate spans for lazy RPCs

}

private final class OtelDecoratedBlobWriteSession implements BlobWriteSession {

private final BlobWriteSession delegate;
private final Span sessionSpan;
private final Tracer tracer;

public OtelDecoratedBlobWriteSession(BlobWriteSession delegate) {
public OtelDecoratedBlobWriteSession(BlobWriteSession delegate, Span sessionSpan) {
this.delegate = delegate;
this.sessionSpan = sessionSpan;
this.tracer =
TracerDecorator.decorate(
Context.current(),
Expand Down Expand Up @@ -1696,13 +1718,16 @@ private class OtelDecoratingWritableByteChannel implements WritableByteChannel {
private final WritableByteChannel delegate;
private final Span openSpan;

private Scope scope;

private OtelDecoratingWritableByteChannel(WritableByteChannel delegate, Span openSpan) {
this.delegate = delegate;
this.openSpan = openSpan;
}

@Override
public int write(ByteBuffer src) throws IOException {
setScope();
return delegate.write(src);
}

Expand All @@ -1713,16 +1738,34 @@ public boolean isOpen() {

@Override
public void close() throws IOException {
setScope();
try {
delegate.close();
} catch (IOException | RuntimeException e) {
openSpan.recordException(e);
openSpan.setStatus(StatusCode.ERROR, e.getClass().getSimpleName());
sessionSpan.recordException(e);
sessionSpan.setStatus(StatusCode.ERROR, e.getClass().getSimpleName());
throw e;
} finally {
openSpan.end();
sessionSpan.end();
clearScope();
}
}

private void clearScope() {
try (Scope ignore = scope) {
scope = null;
}
}

public void setScope() {
if (scope != null) {
clearScope();
}
scope = openSpan.makeCurrent();
}
}
}

Expand All @@ -1731,6 +1774,8 @@ static final class OtelDecoratedWriteChannel implements WriteChannel {
@VisibleForTesting final WriteChannel delegate;
private final Span openSpan;

private Scope scope;

private OtelDecoratedWriteChannel(WriteChannel delegate, Span openSpan) {
this.delegate = delegate;
this.openSpan = openSpan;
Expand All @@ -1748,6 +1793,7 @@ public RestorableState<WriteChannel> capture() {

@Override
public int write(ByteBuffer src) throws IOException {
setScope();
return delegate.write(src);
}

Expand All @@ -1758,6 +1804,7 @@ public boolean isOpen() {

@Override
public void close() throws IOException {
setScope();
try {
delegate.close();
} catch (IOException | RuntimeException e) {
Expand All @@ -1766,7 +1813,21 @@ public void close() throws IOException {
throw e;
} finally {
openSpan.end();
clearScope();
}
}

private void clearScope() {
try (Scope ignore = scope) {
scope = null;
}
}

public void setScope() {
if (scope != null) {
clearScope();
}
scope = openSpan.makeCurrent();
}
}

Expand Down Expand Up @@ -1962,7 +2023,7 @@ public BlobInfo getBlobInfo() {
public <Projection> Projection readAs(ReadProjectionConfig<Projection> config) {
Span readRangeSpan =
tracer
.spanBuilder("readAs")
.spanBuilder(BLOB_READ_SESSION + "/readAs")
.setAttribute("gsutil.uri", id.toGsUtilUriWithGeneration())
.setParent(blobReadSessionContext)
.startSpan();
Expand Down Expand Up @@ -2145,6 +2206,8 @@ private final class OtelDecoratingAppendableUploadWriteableByteChannel
private final AppendableUploadWriteableByteChannel delegate;
private final Span openSpan;

private volatile Scope scope;

private OtelDecoratingAppendableUploadWriteableByteChannel(
AppendableUploadWriteableByteChannel delegate, Span openSpan) {
this.delegate = delegate;
Expand All @@ -2165,6 +2228,7 @@ public void finalizeAndClose() throws IOException {
} finally {
openSpan.end();
uploadSpan.end();
clearScope();
}
}

Expand All @@ -2182,12 +2246,14 @@ public void closeWithoutFinalizing() throws IOException {
} finally {
openSpan.end();
uploadSpan.end();
clearScope();
}
}

@Override
@BetaApi
public void close() throws IOException {
setScope();
try {
delegate.close();
} catch (IOException | RuntimeException e) {
Expand All @@ -2199,18 +2265,33 @@ public void close() throws IOException {
} finally {
openSpan.end();
uploadSpan.end();
clearScope();
}
}

@Override
public int write(ByteBuffer src) throws IOException {
setScope();
return delegate.write(src);
}

@Override
public boolean isOpen() {
return delegate.isOpen();
}

private void clearScope() {
try (Scope ignore = scope) {
scope = null;
}
}

public void setScope() {
if (scope != null) {
clearScope();
}
scope = openSpan.makeCurrent();
}
}
}
}
Loading

Back | FazBrowse Home | New Git URL