aboutsummaryrefslogtreecommitdiff
path: root/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java
blob: aa2f055a02f32a97bb4df86616846d74d8f6ce55 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
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()));
  }
}