nl.topicus.jdbc.shaded.com.google.cloud.http.CensusHttpModule Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of spanner-jdbc Show documentation
JDBC Driver for Google Cloud Spanner
There is a newer version: 1.1.6
/*
 * Copyright 2018 Google LLC
 *
 * 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 nl.topicus.jdbc.shaded.com.google.cloud.http;

import static nl.topicus.jdbc.shaded.com.google.common.base.Preconditions.checkNotNull;
import nl.topicus.jdbc.shaded.com.google.api.client.http.HttpExecuteInterceptor;
import nl.topicus.jdbc.shaded.com.google.api.client.http.HttpHeaders;
import nl.topicus.jdbc.shaded.com.google.api.client.http.HttpRequest;
import nl.topicus.jdbc.shaded.com.google.api.client.http.HttpRequestInitializer;
import nl.topicus.jdbc.shaded.com.google.common.annotations.VisibleForTesting;
import nl.topicus.jdbc.shaded.io.opencensus.contrib.http.util.HttpPropagationUtil;
import nl.topicus.jdbc.shaded.io.opencensus.trace.SpanContext;
import nl.topicus.jdbc.shaded.io.opencensus.trace.Tracer;
import nl.topicus.jdbc.shaded.io.opencensus.trace.propagation.TextFormat;
import java.io.IOException;
import nl.topicus.jdbc.shaded.javax.annotation.Nullable;

/**
 * Provides utilities for Open Census to record http stats/trace information and propagate tracing
 * context.
 */
public final class CensusHttpModule {

  /**
   * OpenCensus tracing component.
   * When no OpenCensus implementation is provided, it will return a no-op tracer.
   */
  private final Tracer tracer;

  /**
   * {@link TextFormat} used in tracing context propagation.
   */
  @Nullable
  private final TextFormat propagationTextFormat;

  /**
   * {@link TextFormat.Setter} for {@link #propagationTextFormat}.
   */
  @Nullable
  private final TextFormat.Setter propagationTextFormatSetter;

  /**
   * Whether spans are stored locally.
   */
  private final boolean isRecordEvents;

  /**
   * Default HTTP propagation text formatter.
   */
  @VisibleForTesting
  final static class DefaultPropagationTextFormatSetter extends TextFormat.Setter {
    static final TextFormat.Setter INSTANCE = new DefaultPropagationTextFormatSetter();

    @Override
    public void put(HttpHeaders carrier, String key, String value) {
      carrier.set(key, value);
    }
  }

  /**
   * An {@link HttpExecuteInterceptor} implementation to inject HTTP request and add getContext
   * information before it is executed.
   */
  @VisibleForTesting
  final class CensusHttpExecuteInterceptor implements HttpExecuteInterceptor {
    @Nullable
    HttpExecuteInterceptor interceptor;

    CensusHttpExecuteInterceptor(HttpExecuteInterceptor interceptor) {
      this.interceptor = interceptor;
    }

    @Override
    public void intercept(HttpRequest request) throws IOException {
      checkNotNull(request);
      if (this.interceptor != null) {
        this.interceptor.intercept(request);
      }
      if (propagationTextFormat != null && propagationTextFormatSetter != null) {
        SpanContext spanContext = tracer.getCurrentSpan().getContext();
        if (!SpanContext.INVALID.equals(spanContext)) {
          propagationTextFormat.inject(
              spanContext, request.getHeaders(), propagationTextFormatSetter);
        }
      }
    }
  }

  /**
   * An {@link HttpRequestInitializer} implementation to set {@link CensusHttpExecuteInterceptor} as
   * interceptor.
   */
  @VisibleForTesting
  final class CensusHttpRequestInitializer implements HttpRequestInitializer {
    @Nullable
    HttpRequestInitializer initializer;

    CensusHttpRequestInitializer(HttpRequestInitializer initializer) {
      this.initializer = initializer;
    }

    @Override
    public void initialize(HttpRequest request) throws IOException {
      checkNotNull(request);
      if (this.initializer != null) {
        this.initializer.initialize(request);
      }
      request.setInterceptor(new CensusHttpExecuteInterceptor(request.getInterceptor()));
    }
  }

  /**
   * Creates a {@link CensusHttpModule} with given parameters.
   *
   * @param tracer the OpenCensus {@code Tracer}.
   * @param isRecordEvents whether spans are stored locally.
   */
  public CensusHttpModule(Tracer tracer, boolean isRecordEvents) {
    checkNotNull(tracer, "tracer");
    this.tracer = tracer;
    this.isRecordEvents = isRecordEvents;
    this.propagationTextFormat = HttpPropagationUtil.getCloudTraceFormat();
    this.propagationTextFormatSetter = DefaultPropagationTextFormatSetter.INSTANCE;
  }

  /**
   * Returns the tracing component of OpenCensus.
   *
   * @return the tracing component of OpenCensus.
   */
  public Tracer getTracer() {
    return tracer;
  }

  /**
   * Returns whether spans are stored locally.
   *
   * @return whether spans are stored locally.
   */
  public boolean isRecordEvents() {
    return isRecordEvents;
  }

  /**
   * Returns the {@link HttpExecuteInterceptor} used when initializing the {@link HttpRequest}.
   *
   * @param initializer the original initializer which will be executed before this initializer.
   * @return the {@code HttpExecuteInterceptor}.
   */
  public HttpRequestInitializer getHttpRequestInitializer(HttpRequestInitializer initializer) {
    return new CensusHttpRequestInitializer(initializer);
  }
}