package org.stagemonitor.tracing.mdc;

import org.slf4j.MDC;
import org.stagemonitor.core.CorePlugin;
import org.stagemonitor.core.MeasurementSession;
import org.stagemonitor.core.Stagemonitor;
import org.stagemonitor.tracing.B3HeaderFormat;
import org.stagemonitor.tracing.TracingPlugin;
import org.stagemonitor.tracing.wrapper.SpanWrapper;
import org.stagemonitor.tracing.wrapper.StatelessSpanEventListener;

/**
 * This class adds the {@link MDC} properties requestId, application, host and instance.
 * <p>
 * If you are using logback or log4j, you can append this to your pattern to append the properties to each log entry:
 * <code>trace:[%X{traceId}] span:[%X{spanId}] A:[%X{application}] H:[%X{host}] I:[%X{instance}]</code>
 */
public class MDCSpanEventListener extends StatelessSpanEventListener {

	private final CorePlugin corePlugin;
	private final TracingPlugin tracingPlugin;

	public MDCSpanEventListener(CorePlugin corePlugin, TracingPlugin tracingPlugin) {
		this.corePlugin = corePlugin;
		this.tracingPlugin = tracingPlugin;
	}

	@Override
	public void onStart(SpanWrapper spanWrapper) {
		if (corePlugin.isStagemonitorActive()) {
			final MeasurementSession measurementSession = corePlugin.getMeasurementSession();
			if (measurementSession != null) {
				addToMdcIfNotNull("application", measurementSession.getApplicationName());
				addToMdcIfNotNull("host", measurementSession.getHostName());
				addToMdcIfNotNull("instance", measurementSession.getInstanceName());
			}

			// don't store the context in MDC if stagemonitor is not active
			// so that thread pools that get created on startup don't inherit the ids
			if (Stagemonitor.isStarted()) {
				tracingPlugin.getTracer().inject(spanWrapper.context(), B3HeaderFormat.INSTANCE, new B3HeaderFormat.B3InjectAdapter() {
					@Override
					public void setParentId(String value) {
						addToMdcIfNotNull("parentId", value);
					}

					@Override
					public void setSpanId(String value) {
						addToMdcIfNotNull("spanId", value);
					}

					@Override
					public void setTraceId(String value) {
						addToMdcIfNotNull("traceId", value);
					}
				});
			}
		}
	}

	@Override
	public void onFinish(SpanWrapper spanWrapper, String operationName, long durationNanos) {
		// the other keys are not span-scoped hence not removed here
		MDC.remove("traceId");
		MDC.remove("spanId");
		MDC.remove("parentId");
	}

	private void addToMdcIfNotNull(String key, String value) {
		if (value != null) {
			MDC.put(key, value);
		}
	}
}

@@ -802,16 +802,271 @@
 ener %7B%0A%0A
+%09private static final String TRACE_ID = %22traceId%22;%0A%09private static final String SPAN_ID = %22spanId%22;%0A%09private static final String APPLICATION = %22application%22;%0A%09private static final String HOST = %22host%22;%0A%09private static final String INSTANCE = %22instance%22;%0A%0A
 %09private
@@ -2111,52 +2111,8 @@
 ) %7B%0A
-%09%09%09%09%09%09addToMdcIfNotNull(%22parentId%22, value);%0A
 %09%09%09%09
@@ -2461,68 +2461,52 @@
 %7B%0A%09%09
-// the other keys are not span-scoped hence not removed here
+MDC.remove(TRACE_ID);%0A%09%09MDC.remove(SPAN_ID);
 %0A%09%09M
@@ -2515,25 +2515,27 @@
 .remove(
-%22traceId%22
+APPLICATION
 );%0A%09%09MDC
@@ -2542,24 +2542,20 @@
 .remove(
-%22spanId%22
+HOST
 );%0A%09%09MDC
@@ -2566,18 +2566,16 @@
 ove(
-%22parentId%22
+INSTANCE
 );%0A%09

Remove all MDC entries onFinish