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.OpenTelemetry;
022import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
023import io.opentelemetry.context.propagation.ContextPropagators;
024import io.opentelemetry.sdk.OpenTelemetrySdk;
025import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
026import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule;
027import io.opentelemetry.sdk.trace.SdkTracerProvider;
028import io.opentelemetry.sdk.trace.data.SpanData;
029import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
030import java.util.List;
031import org.junit.rules.ExternalResource;
032
033/**
034 * <p>
035 * Like {@link OpenTelemetryRule}, except modeled after the junit5 implementation
036 * {@code OpenTelemetryExtension}. Use this class when you need to make asserts on {@link SpanData}
037 * created on a MiniCluster. Make sure this rule initialized before the MiniCluster so that it can
038 * register its instance of {@link OpenTelemetry} as the global instance before any server-side
039 * component can call {@link TraceUtil#getGlobalTracer()}.
040 * </p>
041 * <p>
042 * For example:
043 * </p>
044 *
045 * <pre>
046 * {
047 *   &#64;code
048 *   public class TestMyClass {
049 *     private static final OpenTelemetryClassRule otelClassRule = OpenTelemetryClassRule.create();
050 *     private static final MiniClusterRule miniClusterRule = MiniClusterRule.newBuilder().build();
051 *     protected static final ConnectionRule connectionRule =
052 *       ConnectionRule.createAsyncConnectionRule(miniClusterRule::createAsyncConnection);
053 *
054 *     &#64;ClassRule
055 *     public static final TestRule classRule =
056 *       RuleChain.outerRule(otelClassRule).around(miniClusterRule).around(connectionRule);
057 *
058 *     &#64;Rule
059 *     public final OpenTelemetryTestRule otelTestRule = new OpenTelemetryTestRule(otelClassRule);
060 *
061 *     &#64;Test
062 *     public void myTest() {
063 *       // ...
064 *       // do something that makes spans
065 *       final List<SpanData> spans = otelClassRule.getSpans();
066 *       // make assertions on them
067 *     }
068 *   }
069 * }
070 * </pre>
071 *
072 * @see <a href=
073 *      "https://github.com/open-telemetry/opentelemetry-java/blob/9a330d0/sdk/testing/src/main/java/io/opentelemetry/sdk/testing/junit5/OpenTelemetryExtension.java">junit5/OpenTelemetryExtension.java</a>
074 */
075public final class OpenTelemetryClassRule extends ExternalResource {
076
077  public static OpenTelemetryClassRule create() {
078    InMemorySpanExporter spanExporter = InMemorySpanExporter.create();
079
080    SdkTracerProvider tracerProvider = SdkTracerProvider.builder()
081      .addSpanProcessor(SimpleSpanProcessor.create(spanExporter)).build();
082
083    OpenTelemetrySdk openTelemetry = OpenTelemetrySdk.builder()
084      .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
085      .setTracerProvider(tracerProvider).build();
086
087    return new OpenTelemetryClassRule(openTelemetry, spanExporter);
088  }
089
090  private final OpenTelemetrySdk openTelemetry;
091  private final InMemorySpanExporter spanExporter;
092
093  private OpenTelemetryClassRule(final OpenTelemetrySdk openTelemetry,
094    final InMemorySpanExporter spanExporter) {
095    this.openTelemetry = openTelemetry;
096    this.spanExporter = spanExporter;
097  }
098
099  /** Returns the {@link OpenTelemetry} created by this Rule. */
100  public OpenTelemetry getOpenTelemetry() {
101    return openTelemetry;
102  }
103
104  /** Returns all the exported {@link SpanData} so far. */
105  public List<SpanData> getSpans() {
106    return spanExporter.getFinishedSpanItems();
107  }
108
109  /**
110   * Clears the collected exported {@link SpanData}.
111   */
112  public void clearSpans() {
113    spanExporter.reset();
114  }
115
116  @Override
117  protected void before() throws Throwable {
118    GlobalOpenTelemetry.resetForTest();
119    GlobalOpenTelemetry.set(openTelemetry);
120  }
121
122  @Override
123  protected void after() {
124    GlobalOpenTelemetry.resetForTest();
125  }
126}