package smartthings.dw.asynchttpclient;

import io.netty.handler.codec.http.HttpHeaders;
import org.asynchttpclient.*;
import org.asynchttpclient.filter.FilterContext;
import org.asynchttpclient.filter.FilterException;
import org.asynchttpclient.filter.RequestFilter;
import org.slf4j.MDC;
import smartthings.dw.logging.LoggingContext;

import java.util.Map;

public class CorrelationIdFilter implements RequestFilter {

	@Override
	public <T> FilterContext<T> filter(FilterContext<T> ctx) throws FilterException {

		final Map<String, String> mdc = MDC.getCopyOfContextMap();
		final AsyncHandler<T> asyncHandler = ctx.getAsyncHandler();

		// Wrap AsyncHandler to copy MDC since it executes on a different thread
		FilterContext<T> newContext = new FilterContext.FilterContextBuilder<>(ctx)
				.asyncHandler(new AsyncHandler<T>() {
					@Override
					public void onThrowable(Throwable t) {
						try {
							mdc.forEach(MDC::put);
							asyncHandler.onThrowable(t);
						} finally {
							MDC.clear();
						}
					}

					@Override
					public State onBodyPartReceived(HttpResponseBodyPart bodyPart) throws Exception {
						State result;
						try {
							mdc.forEach(MDC::put);
							result = asyncHandler.onBodyPartReceived(bodyPart);
						} finally {
							MDC.clear();
						}
						return result;
					}

					@Override
					public State onStatusReceived(HttpResponseStatus responseStatus) throws Exception {
						try {
							mdc.forEach(MDC::put);
							return asyncHandler.onStatusReceived(responseStatus);
						} finally {
							MDC.clear();
						}
					}

					@Override
					public State onHeadersReceived(HttpResponseHeaders headers) throws Exception {
						try {
							mdc.forEach(MDC::put);
							return asyncHandler.onHeadersReceived(headers);
						} finally {
							MDC.clear();
						}
					}

					@Override
					public T onCompleted() throws Exception {
						try {
							mdc.forEach(MDC::put);
							return asyncHandler.onCompleted();
						} finally {
							MDC.clear();
						}
					}
				}).build();

		String loggingId = LoggingContext.getLoggingId();

		if (loggingId != null) {
			HttpHeaders headers = newContext.getRequest().getHeaders();
			if (!headers.contains(LoggingContext.CORRELATION_ID_HEADER)) {
				headers.add(LoggingContext.CORRELATION_ID_HEADER, loggingId);
			}
		}
		return newContext;
	}
}

@@ -326,24 +326,54 @@
 ngContext;%0A%0A
+import java.util.Collections;%0A
 import java.
@@ -573,20 +573,16 @@
 %3E mdc = 
-MDC.
 getCopyO
@@ -2377,11 +2377,271 @@
 text;%0A%09%7D
+%0A%0A    /**%0A     * Provides null safe access to MDC Context Map%0A     */%0A%09private Map%3CString, String%3E getCopyOfContextMap() %7B%0A        Map%3CString, String%3E context = MDC.getCopyOfContextMap();%0A        return context == null ? Collections.emptyMap() : context;%0A    %7D
 %0A%7D%0A

Fix NPE when MDC context map is null