diff options
Diffstat (limited to 'api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java')
-rw-r--r-- | api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java | 180 |
1 files changed, 180 insertions, 0 deletions
diff --git a/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java b/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java new file mode 100644 index 00000000..aa2f055a --- /dev/null +++ b/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java @@ -0,0 +1,180 @@ +/* + * Copyright 2017, OpenCensus Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.opencensus.trace; + +import io.grpc.Context; +import io.opencensus.common.Scope; +import io.opencensus.trace.unsafe.ContextUtils; +import java.util.concurrent.Callable; +import javax.annotation.Nullable; + +/** Util methods/functionality to interact with the {@link Span} in the {@link io.grpc.Context}. */ +final class CurrentSpanUtils { + // No instance of this class. + private CurrentSpanUtils() {} + + /** + * Returns The {@link Span} from the current context. + * + * @return The {@code Span} from the current context. + */ + @Nullable + static Span getCurrentSpan() { + return ContextUtils.CONTEXT_SPAN_KEY.get(); + } + + /** + * Enters the scope of code where the given {@link Span} is in the current context, and returns an + * object that represents that scope. The scope is exited when the returned object is closed. + * + * <p>Supports try-with-resource idiom. + * + * @param span The {@code Span} to be set to the current context. + * @param endSpan if {@code true} the returned {@code Scope} will close the {@code Span}. + * @return An object that defines a scope where the given {@code Span} is set to the current + * context. + */ + static Scope withSpan(Span span, boolean endSpan) { + return new ScopeInSpan(span, endSpan); + } + + /** + * Wraps a {@link Runnable} so that it executes with the {@code span} as the current {@code Span}. + * + * @param span the {@code Span} to be set as current. + * @param endSpan if {@code true} the returned {@code Runnable} will close the {@code Span}. + * @param runnable the {@code Runnable} to run in the {@code Span}. + * @return the wrapped {@code Runnable}. + */ + static Runnable withSpan(Span span, boolean endSpan, Runnable runnable) { + return new RunnableInSpan(span, runnable, endSpan); + } + + /** + * Wraps a {@link Callable} so that it executes with the {@code span} as the current {@code Span}. + * + * @param span the {@code Span} to be set as current. + * @param endSpan if {@code true} the returned {@code Runnable} will close the {@code Span}. + * @param callable the {@code Callable} to run in the {@code Span}. + * @return the wrapped {@code Callable}. + */ + static <C> Callable<C> withSpan(Span span, boolean endSpan, Callable<C> callable) { + return new CallableInSpan<C>(span, callable, endSpan); + } + + // Defines an arbitrary scope of code as a traceable operation. Supports try-with-resources idiom. + private static final class ScopeInSpan implements Scope { + private final Context origContext; + private final Span span; + private final boolean endSpan; + + /** + * Constructs a new {@link ScopeInSpan}. + * + * @param span is the {@code Span} to be added to the current {@code io.grpc.Context}. + */ + private ScopeInSpan(Span span, boolean endSpan) { + this.span = span; + this.endSpan = endSpan; + origContext = Context.current().withValue(ContextUtils.CONTEXT_SPAN_KEY, span).attach(); + } + + @Override + public void close() { + Context.current().detach(origContext); + if (endSpan) { + span.end(); + } + } + } + + private static final class RunnableInSpan implements Runnable { + // TODO(bdrutu): Investigate if the extra private visibility increases the generated bytecode. + private final Span span; + private final Runnable runnable; + private final boolean endSpan; + + private RunnableInSpan(Span span, Runnable runnable, boolean endSpan) { + this.span = span; + this.runnable = runnable; + this.endSpan = endSpan; + } + + @Override + public void run() { + Context origContext = + Context.current().withValue(ContextUtils.CONTEXT_SPAN_KEY, span).attach(); + try { + runnable.run(); + } catch (Throwable t) { + setErrorStatus(span, t); + if (t instanceof RuntimeException) { + throw (RuntimeException) t; + } else if (t instanceof Error) { + throw (Error) t; + } + throw new RuntimeException("unexpected", t); + } finally { + Context.current().detach(origContext); + if (endSpan) { + span.end(); + } + } + } + } + + private static final class CallableInSpan<V> implements Callable<V> { + private final Span span; + private final Callable<V> callable; + private final boolean endSpan; + + private CallableInSpan(Span span, Callable<V> callable, boolean endSpan) { + this.span = span; + this.callable = callable; + this.endSpan = endSpan; + } + + @Override + public V call() throws Exception { + Context origContext = + Context.current().withValue(ContextUtils.CONTEXT_SPAN_KEY, span).attach(); + try { + return callable.call(); + } catch (Exception e) { + setErrorStatus(span, e); + throw e; + } catch (Throwable t) { + setErrorStatus(span, t); + if (t instanceof Error) { + throw (Error) t; + } + throw new RuntimeException("unexpected", t); + } finally { + Context.current().detach(origContext); + if (endSpan) { + span.end(); + } + } + } + } + + private static void setErrorStatus(Span span, Throwable t) { + span.setStatus( + Status.UNKNOWN.withDescription( + t.getMessage() == null ? t.getClass().getSimpleName() : t.getMessage())); + } +} |