ExecutorService/ScheduledExecutorServiceでMDC引き回すやつ

DelegatingMDCInheritableExecutorService.java

package com.github.wreulicke.mdc;

import org.slf4j.MDC;

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

public class DelegatingMDCInheritableExecutorService extends AbstractExecutorService {

    private final ExecutorService delegate;

    public DelegatingMDCInheritableExecutorService(ExecutorService delegate) {
        this.delegate = delegate;
    }

    private Runnable decorate(Runnable runnable) {
        Map<String, String> context = MDC.getCopyOfContextMap();
        return () -> {
            Map<String, String> original = MDC.getCopyOfContextMap();

            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            try {
                runnable.run();
            } finally {
                if (original == null) {
                    MDC.clear();
                } else {
                    MDC.setContextMap(original);
                }
            }
        };
    }

    // delegates
    @Override
    public void shutdown() {
        delegate.shutdown();
    }

    @Override
    public List<Runnable> shutdownNow() {
        return delegate.shutdownNow();
    }

    @Override
    public boolean isShutdown() {
        return delegate.isShutdown();
    }

    @Override
    public boolean isTerminated() {
        return delegate.isTerminated();
    }

    @Override
    public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        return delegate.awaitTermination(timeout, unit);
    }

    @Override
    public void execute(Runnable command) {
        delegate.execute(decorate(command));
    }
}

DelegatingMDCInheritableExecutorServiceTest.java

package com.github.wreulicke.mdc;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.slf4j.MDC;

import java.util.Collections;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import static org.junit.jupiter.api.Assertions.*;

class DelegatingMDCInheritableExecutorServiceTest {

    ExecutorService sut = new DelegatingMDCInheritableExecutorService(Executors.newSingleThreadExecutor());

    @BeforeEach
    void setUp() {
        MDC.put("key", "value");
    }

    @AfterEach
    void tearDown() {
        MDC.clear();
    }

    @Test
    void execute() {
        sut.execute(() -> {
            assertEquals("value", MDC.get("key"));
            System.out.println("here");
        });
    }

    @Test
    void submitWithRunnable() throws ExecutionException, InterruptedException {
        sut.submit(() -> {
            assertEquals("value", MDC.get("key"));
        }).get();
    }

    @Test
    void submitWithCallable() throws ExecutionException, InterruptedException {
        String actual = sut.submit(() -> MDC.get("key")).get();
        assertEquals("value", actual);
    }

    @Test
    void invokeAll() throws InterruptedException, ExecutionException {
        Object actual = sut.invokeAll(Collections.singleton(() -> MDC.get("key"))).get(0).get();
        assertEquals("value", actual);
    }

    @Test
    void invokeAny() throws ExecutionException, InterruptedException {
        Object actual = sut.invokeAny(Collections.singleton(() -> MDC.get("key")));
        assertEquals("value", actual);
    }
}

DelegatingMDCInheritableScheduledExecutorService.java

package com.github.wreulicke.mdc;

import org.slf4j.MDC;

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

public class DelegatingMDCInheritableScheduledExecutorService extends AbstractExecutorService implements ScheduledExecutorService {

    private final ScheduledExecutorService delegate;

    public DelegatingMDCInheritableScheduledExecutorService(ScheduledExecutorService delegate) {
        this.delegate = delegate;
    }

    // delegates
    @Override
    public void shutdown() {
        delegate.shutdown();
    }

    @Override
    public List<Runnable> shutdownNow() {
        return delegate.shutdownNow();
    }

    @Override
    public boolean isShutdown() {
        return delegate.isShutdown();
    }

    @Override
    public boolean isTerminated() {
        return delegate.isTerminated();
    }

    @Override
    public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        return delegate.awaitTermination(timeout, unit);
    }

    @Override
    public void execute(Runnable command) {
        delegate.execute(decorate(command));
    }

    private <T> Callable<T> decorate(Callable<T> callable) {
        Map<String, String> context = MDC.getCopyOfContextMap();
        return () -> {
            Map<String, String> original = MDC.getCopyOfContextMap();

            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            try {
                return callable.call();
            } finally {
                if (original == null) {
                    MDC.clear();
                } else {
                    MDC.setContextMap(original);
                }
            }
        };
    }

    private Runnable decorate(Runnable runnable) {
        Map<String, String> context = MDC.getCopyOfContextMap();
        return () -> {
            Map<String, String> original = MDC.getCopyOfContextMap();

            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            try {
                runnable.run();
            } finally {
                if (original == null) {
                    MDC.clear();
                } else {
                    MDC.setContextMap(original);
                }
            }
        };
    }

    @Override
    public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
        return delegate.schedule(decorate(command), delay, unit);
    }

    @Override
    public <V> ScheduledFuture<V> schedule(Callable<V> callable, long delay, TimeUnit unit) {
        return delegate.schedule(decorate(callable), delay, unit);
    }

    @Override
    public ScheduledFuture<?> scheduleAtFixedRate(Runnable command, long initialDelay, long period, TimeUnit unit) {
        return delegate.scheduleAtFixedRate(decorate(command), initialDelay, period, unit);
    }

    @Override
    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable command, long initialDelay, long delay, TimeUnit unit) {
        return delegate.scheduleAtFixedRate(decorate(command), initialDelay, delay, unit);
    }
}

DelegatingMDCInheritableScheduledExecutorServiceTest.java

package com.github.wreulicke.mdc;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.slf4j.MDC;

import java.util.concurrent.*;

import static org.junit.jupiter.api.Assertions.*;

class DelegatingMDCInheritableScheduledExecutorServiceTest {

    ScheduledExecutorService sut = new DelegatingMDCInheritableScheduledExecutorService(Executors.newScheduledThreadPool(1));

    @BeforeEach
    void setUp() {
        MDC.put("key", "value");
    }

    @AfterEach
    void tearDown() {
        MDC.clear();
    }


    @Test
    void schedule() throws ExecutionException, InterruptedException {
        sut.schedule(() -> {
            assertEquals("value", MDC.get("key"));
        }, 0, TimeUnit.SECONDS).get();
    }

    @Test
    void scheduleWithCallable() throws ExecutionException, InterruptedException {
        String actual = sut.schedule(() -> MDC.get("key"), 0, TimeUnit.SECONDS).get();
        assertEquals("value", actual);
    }

    @Test
    void scheduleAtFixedRate() throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(1);
        ScheduledFuture<?> future = sut.scheduleAtFixedRate(() -> {
            assertEquals("value", MDC.get("key"));
            latch.countDown();
        }, 0, 1000, TimeUnit.SECONDS);
        latch.await();
        future.cancel(true);
    }

    @Test
    void scheduleWithFixedDelay() throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(1);
        ScheduledFuture<?> future = sut.scheduleWithFixedDelay(() -> {
            assertEquals("value", MDC.get("key"));
            latch.countDown();
        }, 0, 1000, TimeUnit.SECONDS);
        latch.await();
        future.cancel(true);
    }
}