content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.instrumentation.vaadin;
  
  import static io.opentelemetry.javaagent.instrumentation.vaadin.VaadinSingletons.REQUEST_HANDLER_CONTEXT_KEY;
  import static io.opentelemetry.javaagent.instrumentation.vaadin.VaadinSingletons.SERVICE_CONTEXT_KEY;
  
  import com.vaadin.flow.component.UI;
  import com.vaadin.flow.router.Location;
  import io.opentelemetry.context.Context;
  import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
  import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRoute;
  import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRouteSource;
  import io.opentelemetry.javaagent.bootstrap.servlet.ServletContextPath;
  import javax.annotation.Nullable;
  
  public class VaadinHelper {
    private final Instrumenter<VaadinHandlerRequest, Void> requestHandlerInstrumenter;
    private final Instrumenter<VaadinServiceRequest, Void> serviceInstrumenter;
  
    public VaadinHelper(
        Instrumenter<VaadinHandlerRequest, Void> requestHandlerInstrumenter,
        Instrumenter<VaadinServiceRequest, Void> serviceInstrumenter) {
      this.requestHandlerInstrumenter = requestHandlerInstrumenter;
      this.serviceInstrumenter = serviceInstrumenter;
    }
  
    @Nullable
    public Context startVaadinServiceSpan(VaadinServiceRequest request) {
      Context parentContext = Context.current();
      if (!serviceInstrumenter.shouldStart(parentContext, request)) {
        return null;
      }
  
      return serviceInstrumenter.start(parentContext, request);
    }
  
    public void endVaadinServiceSpan(
        Context context, VaadinServiceRequest request, Throwable throwable) {
      serviceInstrumenter.end(context, request, null, throwable);
  
      HttpServerRoute.update(
          context,
          HttpServerRouteSource.CONTROLLER,
          VaadinHelper::getSpanNameForVaadinServiceContext,
          request);
    }
  
    private static String getSpanNameForVaadinServiceContext(
        Context context, VaadinServiceRequest request) {
      VaadinServiceContext vaadinServiceContext = context.get(SERVICE_CONTEXT_KEY);
      // None of the request handlers processed the request, set span name to main request processing
      // method name that calls request handlers.
      if (!vaadinServiceContext.isRequestHandled()) {
        return request.getSpanName();
      }
  
      // Name of the request handler that handled the request.
      return vaadinServiceContext.getSpanNameCandidate();
    }
  
    @Nullable
    public Context startRequestHandlerSpan(VaadinHandlerRequest request) {
      Context parentContext = Context.current();
      // ignore nested request handlers
      if (parentContext.get(REQUEST_HANDLER_CONTEXT_KEY) != null) {
        return null;
      }
  
      VaadinServiceContext vaadinServiceContext = parentContext.get(SERVICE_CONTEXT_KEY);
      if (vaadinServiceContext != null && !vaadinServiceContext.isRequestHandled()) {
        // We don't really know whether this request handler is going to be the one that processes
        // the request, if it isn't then next handler will also update server span name candidate.
        vaadinServiceContext.setSpanNameCandidate(request.getSpanName());
      }
  
      if (!requestHandlerInstrumenter.shouldStart(parentContext, request)) {
        return null;
      }
  
      return requestHandlerInstrumenter.start(parentContext, request);
    }
  
    public void endRequestHandlerSpan(
        Context context, VaadinHandlerRequest request, Throwable throwable, boolean handled) {
      requestHandlerInstrumenter.end(context, request, null, throwable);
  
      // request handler returns true when it processes the request, if that is the case then
      // mark request as handled
      if (handled) {
        VaadinServiceContext vaadinServiceContext = context.get(SERVICE_CONTEXT_KEY);
        if (vaadinServiceContext != null) {
          vaadinServiceContext.setRequestHandled();
        }
      }
    }
  
    public void updateServerSpanName(UI ui) {
      if (ui != null) {
        Location location = ui.getInternals().getActiveViewLocation();
        updateServerSpanName(location);
      }
    }
  
    public void updateServerSpanName(Location location) {
      Context context = Context.current();
      HttpServerRoute.update(
          context,
          HttpServerRouteSource.NESTED_CONTROLLER,
          (ctx, loc) -> ServletContextPath.prepend(ctx, getSpanNameForLocation(loc)),
          location);
    }
  
    private static String getSpanNameForLocation(Location location) {
      String path = location.getPath();
      if (!path.isEmpty()) {
        path = "/" + path;
      }
      return path;
    }
  }
  