com.undefinedlabs.scope.context.instrumentation.ThreadDistributedContextScopeRule Maven / Gradle / Ivy

Go to download
package com.undefinedlabs.scope.context.instrumentation;

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

import com.undefinedlabs.scope.bootstrap.ContextTrampoline;
import com.undefinedlabs.scope.rules.AbstractScopeAgentRule;
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;

public class ThreadDistributedContextScopeRule extends AbstractScopeAgentRule {

  @Override
  protected ElementMatcher typeMatcher() {
    return hasSuperType(named("java.lang.Thread")).or(named("java.lang.Thread"));
  }

  @Override
  protected Map, String> transformers() {
    final Map, String> transformers = new HashMap<>();
    transformers.put(
        named("start"), ThreadDistributedContextScopeRule.class.getName() + "$ThreadStartAdvice");
    transformers.put(
        named("run"), ThreadDistributedContextScopeRule.class.getName() + "$ThreadRunAdvice");

    return transformers;
  }

  public static class ThreadStartAdvice {
    @Advice.OnMethodEnter
    public static void enter(@Advice.This final Thread thread) {
      ContextTrampoline.saveContextForThread(thread);
    }
  }

  public static class ThreadRunAdvice {
    @Advice.OnMethodEnter
    public static void enter(@Advice.This final Thread thread) {
      ContextTrampoline.attachContextForThread(thread);
    }
  }
}