com.undefinedlabs.scope.coverage.session.distributed.instrumentation.NettyChannelDistributedContextScopeRule Maven / Gradle / Ivy

package com.undefinedlabs.scope.coverage.session.distributed.instrumentation;

import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
import static net.bytebuddy.matcher.ElementMatchers.named;

import com.undefinedlabs.scope.context.ImmutableContext;
import com.undefinedlabs.scope.coverage.CoverageSession;
import com.undefinedlabs.scope.coverage.rules.AbstractScopeCoverageAgentRule;
import com.undefinedlabs.scope.coverage.session.distributed.DistributedContextCoverageSessionUtils;
import com.undefinedlabs.scope.logger.ScopeLogger;
import com.undefinedlabs.scope.logger.ScopeLoggerResolver;
import java.util.HashMap;
import java.util.Map;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import org.apache.commons.lang3.reflect.MethodUtils;

public class NettyChannelDistributedContextScopeRule extends AbstractScopeCoverageAgentRule {

  public static final ScopeLogger LOGGER = ScopeLoggerResolver.INSTANCE.get();

  @Override
  protected ElementMatcher typeMatcher() {
    return hasSuperType(named("io.netty.channel.Channel"));
  }

  @Override
  protected Map, String> transformers() {
    final Map, String> transformers = new HashMap<>();
    transformers.put(
        named("read"),
        NettyChannelDistributedContextScopeRule.class.getName() + "$ChannelReadAdvice");
    return transformers;
  }

  public static class ChannelReadAdvice {
    @Advice.OnMethodEnter
    public static void enter(@Advice.This final Object channelObj) {
      final CoverageSession activeSession =
          DistributedContextCoverageSessionUtils.getActiveSessionTls(ImmutableContext.current());

      try {
        final Object channelIdObj = MethodUtils.invokeMethod(channelObj, "id");

        final Class attributeKeyClass = Class.forName("io.netty.util.AttributeKey");
        final Object attributeKeyObj =
            MethodUtils.invokeStaticMethod(
                attributeKeyClass, "valueOf", CoverageSession.ACTIVE_SESSION_ON_REGISTRATION);
        final Object attributeChannelObj =
            MethodUtils.invokeMethod(channelObj, "attr", attributeKeyObj);

        final String sessionIdOnRegistration =
            (String) MethodUtils.invokeMethod(attributeChannelObj, "get");
        DistributedContextCoverageSessionUtils.setSessionIdOnChannelRegistration(
            ImmutableContext.current(), sessionIdOnRegistration);

        LOGGER.trace(
            "--- NIO CHANNEL [READ]: ChannelId: ["
                + channelIdObj
                + "] ActiveSession: ["
                + (activeSession != null ? activeSession.getSessionId() : null)
                + "]. ActiveSession on Registration ["
                + sessionIdOnRegistration
                + "].  "
                + Thread.currentThread());
      } catch (final Throwable e) {
        LOGGER.error("--- NIO CHANNEL [READ] Throwable: " + e);
      }
    }
  }
}