package org.pmiops.workbench.interceptors;

import io.opencensus.common.Scope;
import io.opencensus.exporter.trace.stackdriver.StackdriverTraceConfiguration;
import io.opencensus.exporter.trace.stackdriver.StackdriverTraceExporter;
import io.opencensus.trace.SpanBuilder;
import io.opencensus.trace.Tracer;
import io.opencensus.trace.Tracing;
import io.opencensus.trace.samplers.Samplers;
import java.io.IOException;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.inject.Provider;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import org.pmiops.workbench.config.WorkbenchConfig;
import org.pmiops.workbench.firecloud.FirecloudApiClientTracer;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.ModelAndView;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

// Interceptor to create a trace of the lifecycle of api calls.
@Service
public class TracingInterceptor extends HandlerInterceptorAdapter {
  private static final Tracer tracer = Tracing.getTracer();
  private static final Logger log = Logger.getLogger(FirecloudApiClientTracer.class.getName());
  private static final String TRACE_ATTRIBUTE_KEY = "Tracing Span";

  private final Provider<WorkbenchConfig> workbenchConfigProvider;

  @Autowired
  public TracingInterceptor(Provider<WorkbenchConfig> workbenchConfigProvider) {
    this.workbenchConfigProvider = workbenchConfigProvider;
    try {
      StackdriverTraceExporter.createAndRegister(StackdriverTraceConfiguration.builder().build());
    } catch (IOException e) {
      log.log(Level.WARNING, "Failed to setup tracing", e);
    }
  }

  /**
   * @param handler The Swagger-generated ApiController. It contains our handler as a private
   *     delegate.
   */
  @Override
  public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
      throws Exception {

    SpanBuilder requestSpanBuilder =
        tracer.spanBuilder(
            workbenchConfigProvider.get().server.shortName + request.getRequestURI());

    if (workbenchConfigProvider.get().server.traceAllRequests) {
      requestSpanBuilder.setSampler(Samplers.alwaysSample());
    }

    Scope requestSpan = requestSpanBuilder.startScopedSpan();
    request.setAttribute(TRACE_ATTRIBUTE_KEY, requestSpan);
    return true;
  }

  @Override
  public void postHandle(
      HttpServletRequest request,
      HttpServletResponse response,
      Object handler,
      ModelAndView modelAndView)
      throws Exception {
    ((Scope) request.getAttribute(TRACE_ATTRIBUTE_KEY)).close();
  }
}

@@ -828,61 +828,8 @@
 ce;%0A
-import org.springframework.web.servlet.ModelAndView;%0A
 impo
@@ -2084,16 +2084,92 @@
 hortName
+%0A                + %22/%22%0A                + request.getMethod()%0A               
  + reque
@@ -2498,18 +2498,23 @@
 oid 
-postHandle
+afterCompletion
 (%0A  
@@ -2544,22 +2544,16 @@
 request,
-%0A     
  HttpSer
@@ -2574,22 +2574,16 @@
 esponse,
-%0A     
  Object 
@@ -2594,40 +2594,21 @@
 ler,
-%0A      ModelAndView modelAndView
+ Exception ex
 )%0A  

Add request type to trace (#2570)