content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.instrumentation.spring.webmvc.v5_3;
  
  import static io.opentelemetry.instrumentation.api.semconv.http.HttpServerRouteSource.CONTROLLER;
  import static java.util.Objects.requireNonNull;
  
  import io.opentelemetry.context.Context;
  import io.opentelemetry.context.Scope;
  import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
  import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRoute;
  import java.io.IOException;
  import java.util.concurrent.atomic.AtomicBoolean;
  import javax.servlet.AsyncContext;
  import javax.servlet.AsyncEvent;
  import javax.servlet.AsyncListener;
  import javax.servlet.FilterChain;
  import javax.servlet.ServletException;
  import javax.servlet.ServletRequest;
  import javax.servlet.ServletResponse;
  import javax.servlet.http.HttpServletRequest;
  import javax.servlet.http.HttpServletRequestWrapper;
  import javax.servlet.http.HttpServletResponse;
  import org.springframework.core.Ordered;
  import org.springframework.web.filter.OncePerRequestFilter;
  
  final class WebMvcTelemetryProducingFilter extends OncePerRequestFilter implements Ordered {
  
    private final Instrumenter<HttpServletRequest, HttpServletResponse> instrumenter;
    private final HttpRouteSupport httpRouteSupport = new HttpRouteSupport();
  
    WebMvcTelemetryProducingFilter(
        Instrumenter<HttpServletRequest, HttpServletResponse> instrumenter) {
      this.instrumenter = instrumenter;
    }
  
    @Override
    public void afterPropertiesSet() {
      // don't do anything, in particular do not call initFilterBean()
    }
  
    @Override
    protected void initFilterBean() {
      // FilterConfig must be non-null at this point
      httpRouteSupport.onFilterInit(requireNonNull(getFilterConfig()));
    }
  
    @Override
    public void doFilterInternal(
        HttpServletRequest request, HttpServletResponse response, FilterChain filterChain)
        throws ServletException, IOException {
  
      Context parentContext = Context.current();
      if (!instrumenter.shouldStart(parentContext, request)) {
        filterChain.doFilter(request, response);
        return;
      }
  
      Context context = instrumenter.start(parentContext, request);
      AsyncAwareHttpServletRequest asyncAwareRequest =
          new AsyncAwareHttpServletRequest(request, response, context);
      Throwable error = null;
      try (Scope ignored = context.makeCurrent()) {
        filterChain.doFilter(asyncAwareRequest, response);
      } catch (Throwable t) {
        error = t;
        throw t;
      } finally {
        if (httpRouteSupport.hasMappings()) {
          HttpServerRoute.update(context, CONTROLLER, httpRouteSupport::getHttpRoute, request);
        }
        if (error != null || asyncAwareRequest.isNotAsync()) {
          instrumenter.end(context, request, response, error);
        }
      }
    }
  
    @Override
    public void destroy() {}
  
    @Override
    public int getOrder() {
      // Run after all HIGHEST_PRECEDENCE items
      return Ordered.HIGHEST_PRECEDENCE + 1;
    }
  
    private class AsyncAwareHttpServletRequest extends HttpServletRequestWrapper {
      private final HttpServletRequest request;
      private final HttpServletResponse response;
      private final Context context;
      private final AtomicBoolean listenerAttached = new AtomicBoolean();
  
      AsyncAwareHttpServletRequest(
          HttpServletRequest request, HttpServletResponse response, Context context) {
        super(request);
        this.request = request;
        this.response = response;
        this.context = context;
      }
  
      @Override
      public AsyncContext startAsync() {
        AsyncContext asyncContext = super.startAsync();
        attachListener(asyncContext);
        return asyncContext;
      }
  
      @Override
      public AsyncContext startAsync(ServletRequest servletRequest, ServletResponse servletResponse) {
        AsyncContext asyncContext = super.startAsync(servletRequest, servletResponse);
        attachListener(asyncContext);
        return asyncContext;
      }
  
      private void attachListener(AsyncContext asyncContext) {
        if (!listenerAttached.compareAndSet(false, true)) {
          return;
        }
  
        asyncContext.addListener(
            new AsyncRequestCompletionListener(request, response, context), request, response);
      }
  
      boolean isNotAsync() {
        return !listenerAttached.get();
      }
    }
  
    private class AsyncRequestCompletionListener implements AsyncListener {
      private final HttpServletRequest request;
      private final HttpServletResponse response;
      private final Context context;
      private final AtomicBoolean responseHandled = new AtomicBoolean();
  
      AsyncRequestCompletionListener(
          HttpServletRequest request, HttpServletResponse response, Context context) {
        this.request = request;
        this.response = response;
        this.context = context;
      }
  
      @Override
      public void onComplete(AsyncEvent asyncEvent) {
        if (responseHandled.compareAndSet(false, true)) {
          instrumenter.end(context, request, response, null);
        }
      }
  
      @Override
      public void onTimeout(AsyncEvent asyncEvent) {
        if (responseHandled.compareAndSet(false, true)) {
          instrumenter.end(context, request, response, null);
        }
      }
  
      @Override
      public void onError(AsyncEvent asyncEvent) {
        if (responseHandled.compareAndSet(false, true)) {
          instrumenter.end(context, request, response, asyncEvent.getThrowable());
        }
      }
  
      @Override
      public void onStartAsync(AsyncEvent asyncEvent) {
        asyncEvent
            .getAsyncContext()
            .addListener(this, asyncEvent.getSuppliedRequest(), asyncEvent.getSuppliedResponse());
      }
    }
  }
  