opentracing activation hook

Activation.java

package io.opentracing.activation;

import java.util.List;
import java.util.Map;
import org.slf4j.MDC;


public class ActivationTest {

  interface ActiveSpanManager {
    ActivationContext prepare();
  }

  interface ActivationContext extends AutoCloseable {
    // returns this for use with try-with-resource
    ActivationContext start();

    void close();
  }

  // object returned from onPrepare will be passed to onStart and onClose
  interface ActivationHook<T> {
    T onPrepare();
    void onStart(T context);
    void onClose(T context);
  }

  static class InstrumentedRunnable implements Runnable {
    private final ActivationContext context;
    private Runnable wrappedRunnable;

    public InstrumentedRunnable(ActiveSpanManager asm, Runnable wrappedRunnable) {
      context = asm.prepare();
      this.wrappedRunnable = wrappedRunnable;
    }

    @Override
    public void run() {
      try(context.start()) {
        wrappedRunnable.run();
      }
    }
  }

  static class TracerImplementation implements ActiveSpanManager {
    private List<ActivationHook> hooks;

    public ActivationContext prepare() {
      return new CustomActivationContext(hooks);
    }

    static class CustomActivationContext implements ActivationContext {
      private final Object[] state;
      private List<ActivationHook> hooks;

      public CustomActivationContext(List<ActivationHook> hooks) {
        this.hooks = hooks;
        this.state = new Object[hooks.size()];
        for(int i = 0; i < hooks.size(); i++) {
          state[i] = hooks.get(i).onPrepare();
        }
      }

      @Override public ActivationContext start() {
        for(int i = 0; i < hooks.size(); i++) {
          hooks.get(i).onStart(state[i]);
        }
        return this;
      }

      @Override public void close() {
        for(int i = 0; i < hooks.size(); i++) {
          hooks.get(i).onClose(state[i]);
        }
      }
    }
  }


  static class MdcHookExample implements ActivationHook<Map<String,String>> {
    @Override public Map<String,String> onPrepare() {
      return MDC.getCopyOfContextMap();
    }

    @Override public void onStart(Map<String,String> context) {
      MDC.setContextMap(context);
    }

    @Override public void onClose(Map<String,String> context) {
      MDC.clear();
      // real impl would probably restore MDC state from before onStart
    }

    
    static class ThreadLocalHookExample implements ActivationHook<String> {
      ThreadLocal<String> state;
      
      public ThreadLocalHookExample(ThreadLocal<String> state) {
        this.state = state;
      }

      @Override public String onPrepare() {
        return state.get();
      }

      @Override public void onStart(String value) {
        state.set(value);
      }

      @Override public void onClose(String value) {
        state.remove();
      }
    }
  }
}