SpanWrapper.java

package org.stagemonitor.tracing.wrapper;

import com.fasterxml.jackson.annotation.JsonValue;

import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import io.opentracing.Span;
import io.opentracing.SpanContext;

/**
 * The purpose of this wrapper is to call the registered {@link SpanEventListener}s before certain methods of the
 * {@link #delegate} {@link Span} are invoked.
 */
public class SpanWrapper implements Span {

	/**
	 * The actual span to be invoked
	 */
	private Span delegate;
	private String operationName;
	private final long startTimestampNanos;
	private final long startTimestampMillis;
	private final List<SpanEventListener> spanEventListeners;

	public SpanWrapper(Span delegate, String operationName, long startTimestampNanos, long startTimestampMillis, List<SpanEventListener> spanEventListeners) {
		this.delegate = delegate;
		this.operationName = operationName;
		this.startTimestampNanos = startTimestampNanos;
		this.startTimestampMillis = startTimestampMillis;
		this.spanEventListeners = spanEventListeners;
	}

	public SpanContext context() {
		return delegate.context();
	}

	public void close() {
		final long durationNanos = System.nanoTime() - startTimestampNanos;
		for (SpanEventListener spanEventListener : spanEventListeners) {
			spanEventListener.onFinish(this, operationName, durationNanos);
		}
		delegate.close();
	}

	public void finish() {
		final long durationNanos = System.nanoTime() - startTimestampNanos;
		for (SpanEventListener spanEventListener : spanEventListeners) {
			spanEventListener.onFinish(this, operationName, durationNanos);
		}
		delegate.finish();
	}

	public void finish(long finishMicros) {
		final long durationNanos = TimeUnit.MICROSECONDS.toNanos(finishMicros) - startTimestampNanos;
		for (SpanEventListener spanEventListener : spanEventListeners) {
			spanEventListener.onFinish(this, operationName, durationNanos);
		}
		delegate.finish(finishMicros);
	}

	public Span setTag(String key, String value) {
		for (SpanEventListener spanEventListener : spanEventListeners) {
			value = spanEventListener.onSetTag(key, value);
		}
		if (value != null) {
			delegate = delegate.setTag(key, value);
		}
		return this;
	}

	public Span setTag(String key, boolean value) {
		for (SpanEventListener spanEventListener : spanEventListeners) {
			value = spanEventListener.onSetTag(key, value);
		}
		delegate = delegate.setTag(key, value);
		return this;
	}

	public Span setTag(String key, Number value) {
		for (SpanEventListener spanEventListener : spanEventListeners) {
			value = spanEventListener.onSetTag(key, value);
		}
		if (value != null) {
			delegate = delegate.setTag(key, value);
		}
		return this;
	}

	@Override
	public Span log(Map<String, ?> fields) {
		delegate = delegate.log(fields);
		return this;
	}

	@Override
	public Span log(long timestampMicroseconds, Map<String, ?> fields) {
		delegate = delegate.log(timestampMicroseconds, fields);
		return this;
	}

	@Override
	public Span log(String event) {
		delegate = delegate.log(event);
		return this;
	}

	@Override
	public Span log(long timestampMicroseconds, String event) {
		delegate = delegate.log(timestampMicroseconds, event);
		return this;
	}

	public Span log(String eventName, Object payload) {
		delegate = delegate.log(eventName, payload);
		return this;
	}

	public Span log(long timestampMicroseconds, String eventName, Object payload) {
		delegate = delegate.log(timestampMicroseconds, eventName, payload);
		return this;
	}

	public Span setBaggageItem(String key, String value) {
		delegate = delegate.setBaggageItem(key, value);
		return this;
	}

	public String getBaggageItem(String key) {
		return delegate.getBaggageItem(key);
	}

	public Span setOperationName(String operationName) {
		this.operationName = operationName;
		delegate = delegate.setOperationName(operationName);
		return this;
	}

	@JsonValue
	public Span getDelegate() {
		return delegate;
	}

	public String getOperationName() {
		return operationName;
	}

	public long getStartTimestampMillis() {
		return startTimestampMillis;
	}

	public <T extends Span> T unwrap(Class<T> delegateClass) {
		if (delegateClass.isInstance(delegate)) {
			return (T) delegate;
		} else if (delegate instanceof SpanWrapper) {
			return ((SpanWrapper) delegate).unwrap(delegateClass);
		} else {
			return null;
		}
	}
}