DiagnosticsContext.java

import org.slf4j.MDC;

import java.util.HashMap;
import java.util.Map;
import java.util.function.*;

public final class DiagnosticsContext {

    private final Map<String, String> context = new HashMap<>();

    public DiagnosticsContext put(final String key, final Object value) {
        if (key == null) {
            throw new IllegalArgumentException("Illegal argument key: null");
        }

        if (value != null) {
            context.put(key, value.toString());
        }
        return this;
    }

    public void wrap(Runnable block) {
        populate();
        try {
            block.run();
        } finally {
            clean();
        }
    }

    public <T> Consumer<T> wrap(Consumer<T> block) {
        return t -> {
            populate();
            try {
                block.accept(t);
            } finally {
                clean();
            }
        };
    }
    public <T, U> BiConsumer<T, U> wrap(BiConsumer<T, U> block) {
        return (t, u) -> {
            populate();
            try {
                block.accept(t, u);
            } finally {
                clean();
            }
        };
    }

    public <T> Supplier<T> wrap(Supplier<T> block) {
        return () -> {
            populate();
            try {
                return block.get();
            } finally {
                clean();
            }
        };
    }

    public <T, R> Function<T, R> wrap(Function<T, R> block) {
        return t -> {
            populate();
            try {
                return block.apply(t);
            } finally {
                clean();
            }
        };
    }

    public <T, U, R> BiFunction<T, U, R> wrap(BiFunction<T, U, R> block) {
        return (t, u) -> {
            populate();
            try {
                return block.apply(t, u);
            } finally {
                clean();
            }
        };
    }

    private void populate() {
        context.entrySet().forEach(e -> MDC.put(e.getKey(), e.getValue()));
    }

    private void clean() {
        context.keySet().forEach(MDC::remove);
    }

    @Override
    public String toString() {
        return "DiagnosticsContext{" +
                "context=" + context +
                '}';
    }
}