package org.robotninjas.barge;

import com.google.common.collect.Maps;
import org.jetlang.core.BatchExecutorImpl;
import org.jetlang.core.EventReader;
import org.slf4j.MDC;

import java.util.Map;

class BatchExecutor extends BatchExecutorImpl {

  Map contextMap = Maps.newHashMap();

  @Override
  public void execute(EventReader toExecute) {

    Map oldContext = MDC.getCopyOfContextMap();
    MDC.setContextMap(contextMap);

    super.execute(toExecute);

    contextMap = MDC.getCopyOfContextMap();
    MDC.setContextMap(oldContext);

  }

}

@@ -25,16 +25,56 @@
 barge;%0A%0A
+import com.google.common.base.Optional;%0A
 import c
@@ -104,16 +104,16 @@
 t.Maps;%0A
-
 import o
@@ -382,19 +382,29 @@
  %7B%0A%0A    
+Optional%3C
 Map
+%3E
  oldCont
@@ -409,16 +409,38 @@
 ntext = 
+Optional.fromNullable(
 MDC.getC
@@ -448,32 +448,33 @@
 pyOfContextMap()
+)
 ;%0A    MDC.setCon
@@ -570,16 +570,52 @@
 tMap();%0A
+    if (oldContext.isPresent()) %7B%0A  
     MDC.
@@ -638,18 +638,30 @@
 dContext
-);
+.get());%0A    %7D
 %0A%0A  %7D%0A%0A%7D

Handle case where MDC is null