io.opentelemetry.javaagent.instrumentation.undertow.UndertowHelper Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of opentelemetry-javaagent-undertow-1.4 Show documentation
Instrumentation of Java libraries using OpenTelemetry.
There is a newer version: 2.8.0-alpha
Show newest version
/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package io.opentelemetry.javaagent.instrumentation.undertow;

import io.opentelemetry.api.trace.Span;
import io.opentelemetry.context.Context;
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
import io.opentelemetry.javaagent.bootstrap.servlet.AppServerBridge;
import io.opentelemetry.javaagent.bootstrap.undertow.KeyHolder;
import io.opentelemetry.javaagent.bootstrap.undertow.UndertowActiveHandlers;
import io.undertow.server.DefaultResponseListener;
import io.undertow.server.HttpServerExchange;
import io.undertow.util.AttachmentKey;
import javax.annotation.Nullable;

public class UndertowHelper {
  private final Instrumenter instrumenter;

  UndertowHelper(Instrumenter instrumenter) {
    this.instrumenter = instrumenter;
  }

  public boolean shouldStart(Context parentContext, HttpServerExchange exchange) {
    return instrumenter.shouldStart(parentContext, exchange);
  }

  public Context start(Context parentContext, HttpServerExchange exchange) {
    Context context = instrumenter.start(parentContext, exchange);
    attachServerContext(context, exchange);

    return context;
  }

  public void end(Context context, HttpServerExchange exchange, @Nullable Throwable error) {
    if (error == null) {
      error = AppServerBridge.getException(context);
    }

    instrumenter.end(context, exchange, exchange, error);
  }

  public void handlerStarted(Context context) {
    // request was dispatched to a new thread, handler on the original thread
    // may exit before this one so we need to wait for this handler to complete
    // before ending span
    UndertowActiveHandlers.increment(context);
  }

  public void handlerCompleted(Context context, Throwable throwable, HttpServerExchange exchange) {
    // end the span when this is the last handler to complete and exchange has
    // been completed
    if (UndertowActiveHandlers.decrementAndGet(context) == 0) {
      end(context, exchange, throwable);
    }
  }

  public void exchangeCompleted(Context context, HttpServerExchange exchange) {
    // after exchange is completed we can read response status
    // if all handlers have completed we can end the span, if there are running
    // handlers we'll end the span when last handler exits
    if (UndertowActiveHandlers.decrementAndGet(context) == 0) {
      Throwable throwable = exchange.getAttachment(DefaultResponseListener.EXCEPTION);
      end(context, exchange, throwable);
    }
  }

  @SuppressWarnings("unchecked")
  @Nullable
  public Context getServerContext(HttpServerExchange exchange) {
    AttachmentKey contextKey =
        (AttachmentKey) KeyHolder.contextKeys.get(AttachmentKey.class);
    if (contextKey == null) {
      return null;
    }
    return exchange.getAttachment(contextKey);
  }

  @SuppressWarnings("unchecked")
  private static void attachServerContext(Context context, HttpServerExchange exchange) {
    AttachmentKey contextKey =
        (AttachmentKey)
            KeyHolder.contextKeys.computeIfAbsent(
                AttachmentKey.class, key -> AttachmentKey.create(Context.class));
    exchange.putAttachment(contextKey, context);
  }

  public boolean sameTrace(Context currentContext, Context attachedContext) {
    return sameTrace(Span.fromContext(currentContext), Span.fromContext(attachedContext));
  }

  private static boolean sameTrace(Span oneSpan, Span otherSpan) {
    return oneSpan.getSpanContext().getTraceId().equals(otherSpan.getSpanContext().getTraceId());
  }
}