001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.apache.hadoop.hbase.trace;
019
020import io.opentelemetry.api.GlobalOpenTelemetry;
021import io.opentelemetry.api.trace.Span;
022import io.opentelemetry.api.trace.SpanKind;
023import io.opentelemetry.api.trace.StatusCode;
024import io.opentelemetry.api.trace.Tracer;
025import io.opentelemetry.context.Context;
026import io.opentelemetry.context.Scope;
027import java.util.List;
028import java.util.concurrent.Callable;
029import java.util.concurrent.CompletableFuture;
030import java.util.function.Supplier;
031import org.apache.hadoop.hbase.Version;
032import org.apache.hadoop.hbase.util.FutureUtils;
033import org.apache.yetus.audience.InterfaceAudience;
034
035@InterfaceAudience.Private
036public final class TraceUtil {
037
038  private TraceUtil() {
039  }
040
041  public static Tracer getGlobalTracer() {
042    return GlobalOpenTelemetry.getTracer("org.apache.hbase", Version.version);
043  }
044
045  /**
046   * Create a {@link SpanKind#INTERNAL} span.
047   */
048  public static Span createSpan(String name) {
049    return createSpan(name, SpanKind.INTERNAL);
050  }
051
052  /**
053   * Create a span with the given {@code kind}. Notice that, OpenTelemetry only expects one
054   * {@link SpanKind#CLIENT} span and one {@link SpanKind#SERVER} span for a traced request, so use
055   * this with caution when you want to create spans with kind other than {@link SpanKind#INTERNAL}.
056   */
057  private static Span createSpan(String name, SpanKind kind) {
058    return getGlobalTracer().spanBuilder(name).setSpanKind(kind).startSpan();
059  }
060
061  /**
062   * Create a span which parent is from remote, i.e, passed through rpc.
063   * </p>
064   * We will set the kind of the returned span to {@link SpanKind#SERVER}, as this should be the top
065   * most span at server side.
066   */
067  public static Span createRemoteSpan(String name, Context ctx) {
068    return getGlobalTracer().spanBuilder(name).setParent(ctx).setSpanKind(SpanKind.SERVER)
069      .startSpan();
070  }
071
072  /**
073   * Create a span with {@link SpanKind#CLIENT}.
074   */
075  public static Span createClientSpan(String name) {
076    return createSpan(name, SpanKind.CLIENT);
077  }
078
079  /**
080   * Trace an asynchronous operation for a table.
081   */
082  public static <T> CompletableFuture<T> tracedFuture(Supplier<CompletableFuture<T>> action,
083    Supplier<Span> spanSupplier) {
084    Span span = spanSupplier.get();
085    try (Scope ignored = span.makeCurrent()) {
086      CompletableFuture<T> future = action.get();
087      endSpan(future, span);
088      return future;
089    }
090  }
091
092  /**
093   * Trace an asynchronous operation.
094   */
095  public static <T> CompletableFuture<T> tracedFuture(Supplier<CompletableFuture<T>> action,
096    String spanName) {
097    Span span = createSpan(spanName);
098    try (Scope ignored = span.makeCurrent()) {
099      CompletableFuture<T> future = action.get();
100      endSpan(future, span);
101      return future;
102    }
103  }
104
105  /**
106   * Trace an asynchronous operation, and finish the create {@link Span} when all the given
107   * {@code futures} are completed.
108   */
109  public static <T> List<CompletableFuture<T>>
110    tracedFutures(Supplier<List<CompletableFuture<T>>> action, Supplier<Span> spanSupplier) {
111    Span span = spanSupplier.get();
112    try (Scope ignored = span.makeCurrent()) {
113      List<CompletableFuture<T>> futures = action.get();
114      endSpan(CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])), span);
115      return futures;
116    }
117  }
118
119  public static void setError(Span span, Throwable error) {
120    span.recordException(error);
121    span.setStatus(StatusCode.ERROR);
122  }
123
124  /**
125   * Finish the {@code span} when the given {@code future} is completed.
126   */
127  private static void endSpan(CompletableFuture<?> future, Span span) {
128    FutureUtils.addListener(future, (resp, error) -> {
129      if (error != null) {
130        setError(span, error);
131      } else {
132        span.setStatus(StatusCode.OK);
133      }
134      span.end();
135    });
136  }
137
138  /**
139   * Wrap the provided {@code runnable} in a {@link Runnable} that is traced.
140   */
141  public static Runnable tracedRunnable(final Runnable runnable, final String spanName) {
142    return tracedRunnable(runnable, () -> createSpan(spanName));
143  }
144
145  /**
146   * Wrap the provided {@code runnable} in a {@link Runnable} that is traced.
147   */
148  public static Runnable tracedRunnable(final Runnable runnable,
149    final Supplier<Span> spanSupplier) {
150    // N.B. This method name follows the convention of this class, i.e., tracedFuture, rather than
151    // the convention of the OpenTelemetry classes, i.e., Context#wrap.
152    return () -> {
153      final Span span = spanSupplier.get();
154      try (final Scope ignored = span.makeCurrent()) {
155        runnable.run();
156        span.setStatus(StatusCode.OK);
157      } finally {
158        span.end();
159      }
160    };
161  }
162
163  /**
164   * A {@link Runnable} that may also throw.
165   * @param <T> the type of {@link Throwable} that can be produced.
166   */
167  @FunctionalInterface
168  public interface ThrowingRunnable<T extends Throwable> {
169    void run() throws T;
170  }
171
172  /**
173   * Trace the execution of {@code runnable}.
174   */
175  public static <T extends Throwable> void trace(final ThrowingRunnable<T> runnable,
176    final String spanName) throws T {
177    trace(runnable, () -> createSpan(spanName));
178  }
179
180  /**
181   * Trace the execution of {@code runnable}.
182   */
183  public static <T extends Throwable> void trace(final ThrowingRunnable<T> runnable,
184    final Supplier<Span> spanSupplier) throws T {
185    Span span = spanSupplier.get();
186    try (Scope ignored = span.makeCurrent()) {
187      runnable.run();
188      span.setStatus(StatusCode.OK);
189    } catch (Throwable e) {
190      setError(span, e);
191      throw e;
192    } finally {
193      span.end();
194    }
195  }
196
197  /**
198   * A {@link Callable} that may also throw.
199   * @param <R> the result type of method call.
200   * @param <T> the type of {@link Throwable} that can be produced.
201   */
202  @FunctionalInterface
203  public interface ThrowingCallable<R, T extends Throwable> {
204    R call() throws T;
205  }
206
207  public static <R, T extends Throwable> R trace(final ThrowingCallable<R, T> callable,
208    final String spanName) throws T {
209    return trace(callable, () -> createSpan(spanName));
210  }
211
212  public static <R, T extends Throwable> R trace(final ThrowingCallable<R, T> callable,
213    final Supplier<Span> spanSupplier) throws T {
214    Span span = spanSupplier.get();
215    try (Scope ignored = span.makeCurrent()) {
216      final R ret = callable.call();
217      span.setStatus(StatusCode.OK);
218      return ret;
219    } catch (Throwable e) {
220      setError(span, e);
221      throw e;
222    } finally {
223      span.end();
224    }
225  }
226}