Thread Context Management

/**
 * Thread-local Map operations for Mapped Diagnostic Context (MDC)
 */
public final class ThreadContext {
    // Map operations
    /** Put a key-value pair in the thread context */
    public static void put(String key, String value);
    
    /** Put a key-value pair only if the key doesn't exist */
    public static void putIfNull(String key, String value);
    
    /** Put all entries from a map into the thread context */
    public static void putAll(Map<String, String> contextMap);
    
    /** Get value by key from thread context */
    public static String get(String key);
    
    /** Remove a key from the thread context */
    public static void remove(String key);
    
    /** Remove multiple keys from the thread context */
    public static void removeAll(Iterable<String> keys);
    
    /** Clear the entire thread context map */
    public static void clearMap();
    
    /** Get a mutable copy of the current thread context */
    public static Map<String, String> getContext();
    
    /** Get an immutable view of the current thread context */
    public static Map<String, String> getImmutableContext();
    
    /** Check if a key exists in the thread context */
    public static boolean containsKey(String key);
    
    /** Check if the thread context map is empty */
    public static boolean isEmpty();
    
    /** Empty immutable map constant */
    public static final Map<String, String> EMPTY_MAP;
}

private static final Logger logger = LogManager.getLogger();

public void demonstrateThreadContextMap() {
    // Basic key-value operations
    ThreadContext.put("userId", "12345");
    ThreadContext.put("sessionId", "abc-def-ghi");
    ThreadContext.put("requestId", UUID.randomUUID().toString());
    
    logger.info("Processing user request"); // Logs will include context
    
    // Conditional put
    ThreadContext.putIfNull("environment", "production");
    
    // Bulk operations
    Map<String, String> contextData = new HashMap<>();
    contextData.put("correlationId", "corr-123");
    contextData.put("operation", "userLogin");
    ThreadContext.putAll(contextData);
    
    // Reading context
    String userId = ThreadContext.get("userId");
    logger.debug("Current user ID: {}", userId);
    
    // Context existence check
    if (ThreadContext.containsKey("adminMode")) {
        logger.info("Admin mode is active");
    }
    
    // Get full context for processing
    Map<String, String> fullContext = ThreadContext.getContext();
    processWithContext(fullContext);
    
    // Cleanup
    ThreadContext.remove("operation");
    ThreadContext.removeAll(Arrays.asList("userId", "sessionId"));
    ThreadContext.clearMap(); // Clear everything
}

// Web request example
@RestController
public class UserController {
    private static final Logger logger = LogManager.getLogger();
    
    @PostMapping("/users/{userId}/profile")
    public ResponseEntity<UserProfile> updateProfile(
            @PathVariable String userId, 
            @RequestBody ProfileRequest request,
            HttpServletRequest httpRequest) {
        
        // Set up context for entire request
        ThreadContext.put("userId", userId);
        ThreadContext.put("requestId", httpRequest.getHeader("X-Request-ID"));
        ThreadContext.put("endpoint", "/users/{id}/profile");
        ThreadContext.put("method", "POST");
        
        try {
            logger.info("Starting profile update"); // Includes all context
            
            UserProfile profile = userService.updateProfile(userId, request);
            
            logger.info("Profile update completed successfully");
            return ResponseEntity.ok(profile);
            
        } catch (Exception e) {
            logger.error("Profile update failed", e); // Context included in error
            return ResponseEntity.status(500).build();
            
        } finally {
            ThreadContext.clearMap(); // Clean up after request
        }
    }
}

/**
 * Thread-local Stack operations for Nested Diagnostic Context (NDC)
 */
public final class ThreadContext {
    // Stack operations
    /** Push a message onto the thread context stack */
    public static void push(String message);
    
    /** Push a formatted message onto the stack */
    public static void push(String message, Object... args);
    
    /** Pop the top message from the stack and return it */
    public static String pop();
    
    /** Peek at the top message without removing it */
    public static String peek();
    
    /** Clear the entire stack */
    public static void clearStack();
    
    /** Get the current stack depth */
    public static int getDepth();
    
    /** Trim the stack to the specified depth */
    public static void trim(int depth);
    
    /** Create a copy of the current stack */
    public static ContextStack cloneStack();
    
    /** Get an immutable view of the current stack */
    public static ContextStack getImmutableStack();
    
    /** Empty immutable stack constant */
    public static final ThreadContextStack EMPTY_STACK;
    
    /**
     * Stack interface extending Collection<String>
     */
    public interface ContextStack extends Collection<String>, Serializable {
        String pop();
        String peek();
        void push(String message);
        int getDepth();
        List<String> asList();
        void trim(int depth);
        ContextStack copy();
    }
}

private static final Logger logger = LogManager.getLogger();

public void demonstrateThreadContextStack() {
    // Push operation context
    ThreadContext.push("userService");
    logger.info("Starting user operations"); // Stack context included
    
    try {
        ThreadContext.push("validateUser");
        validateUser("12345");
        ThreadContext.pop(); // Remove validateUser
        
        ThreadContext.push("updateProfile");
        updateUserProfile("12345");
        ThreadContext.pop(); // Remove updateProfile
        
    } finally {
        ThreadContext.pop(); // Remove userService
    }
}

// Nested operation tracking
public class OrderProcessor {
    private static final Logger logger = LogManager.getLogger();
    
    public void processOrder(String orderId) {
        ThreadContext.push("processOrder[" + orderId + "]");
        logger.info("Starting order processing");
        
        try {
            validateOrder(orderId);
            processPayment(orderId);
            fulfillOrder(orderId);
            
            logger.info("Order processing completed");
        } finally {
            ThreadContext.pop();
        }
    }
    
    private void validateOrder(String orderId) {
        ThreadContext.push("validateOrder");
        logger.debug("Validating order");
        
        try {
            // Validation logic
            checkInventory(orderId);
            checkCustomer(orderId);
        } finally {
            ThreadContext.pop();
        }
    }
    
    private void checkInventory(String orderId) {
        ThreadContext.push("checkInventory");
        logger.trace("Checking inventory availability");
        
        try {
            // Inventory check logic
        } finally {
            ThreadContext.pop();
        }
    }
}

// Stack manipulation
public void demonstrateStackOperations() {
    ThreadContext.push("operation1");
    ThreadContext.push("operation2");
    ThreadContext.push("operation3");
    
    logger.info("Current depth: {}", ThreadContext.getDepth()); // 3
    logger.info("Top operation: {}", ThreadContext.peek()); // operation3
    
    // Trim to specific depth
    ThreadContext.trim(1); // Keeps only operation1
    
    // Clone stack for processing
    ContextStack stackCopy = ThreadContext.cloneStack();
    processStackCopy(stackCopy);
    
    ThreadContext.clearStack();
}

/**
 * Auto-closeable ThreadContext for automatic cleanup
 */
public class CloseableThreadContext {
    /** Put a key-value pair with automatic cleanup */
    public static Instance put(String key, String value);
    
    /** Put multiple key-value pairs with automatic cleanup */
    public static Instance putAll(Map<String, String> values);
    
    /** Push a message with automatic cleanup */
    public static Instance push(String message);
    
    /** Push multiple messages with automatic cleanup */
    public static Instance pushAll(List<String> messages);
    
    /**
     * Auto-closeable instance for cleanup
     */
    public static class Instance implements AutoCloseable {
        /** Restore previous thread context state */
        @Override
        public void close();
    }
}

private static final Logger logger = LogManager.getLogger();

public void demonstrateCloseableThreadContext() {
    // Automatic cleanup with try-with-resources
    try (CloseableThreadContext.Instance ctc = 
         CloseableThreadContext.put("operation", "userUpdate")) {
        
        logger.info("Starting user update"); // Includes operation context
        updateUser();
        logger.info("User update completed");
        
    } // Context automatically restored here
    
    // Multiple context values
    try (CloseableThreadContext.Instance ctc = 
         CloseableThreadContext.putAll(Map.of(
             "userId", "12345",
             "sessionId", "abc-def",
             "operation", "profileUpdate"))) {
        
        processProfileUpdate();
        
    } // All context automatically cleaned up
    
    // Stack operations with cleanup
    try (CloseableThreadContext.Instance ctc = 
         CloseableThreadContext.push("batchProcessing")) {
        
        processBatch();
        
    } // Stack automatically popped
}

// Web filter example
@Component
public class RequestContextFilter implements Filter {
    
    @Override
    public void doFilter(ServletRequest request, ServletResponse response, 
                        FilterChain chain) throws IOException, ServletException {
        
        HttpServletRequest httpRequest = (HttpServletRequest) request;
        
        // Set up context for entire request with automatic cleanup
        try (CloseableThreadContext.Instance ctc = 
             CloseableThreadContext.putAll(Map.of(
                 "requestId", httpRequest.getHeader("X-Request-ID"),
                 "userAgent", httpRequest.getHeader("User-Agent"),
                 "remoteAddr", httpRequest.getRemoteAddr(),
                 "method", httpRequest.getMethod(),
                 "uri", httpRequest.getRequestURI()))) {
            
            chain.doFilter(request, response);
            
        } // Context automatically cleaned up after request
    }
}

// Nested closeable contexts
public void nestedContextExample() {
    try (CloseableThreadContext.Instance outerCtc = 
         CloseableThreadContext.put("service", "userService")) {
        
        logger.info("Service level logging");
        
        try (CloseableThreadContext.Instance innerCtc = 
             CloseableThreadContext.put("operation", "createUser")) {
            
            logger.info("Operation level logging"); // Both contexts present
            createUser();
            
        } // operation context cleaned up
        
        logger.info("Back to service level"); // Only service context present
        
    } // service context cleaned up
}

// ThreadContext behavior in different threading scenarios
public class ThreadContextBehavior {
    private static final Logger logger = LogManager.getLogger();
    
    public void demonstrateThreadBehavior() {
        // Set context in main thread
        ThreadContext.put("mainThread", "value1");
        logger.info("Main thread logging"); // Includes context
        
        // Context is NOT inherited by new threads
        new Thread(() -> {
            logger.info("New thread logging"); // NO context from main thread
            
            // Each thread has its own context
            ThreadContext.put("workerThread", "value2");
            logger.info("Worker thread with context");
        }).start();
        
        // Main thread still has its context
        logger.info("Back in main thread"); // Still includes mainThread context
    }
    
    // Manual context transfer for async operations
    public CompletableFuture<String> asyncWithContext() {
        // Capture context in current thread
        Map<String, String> contextMap = ThreadContext.getContext();
        ContextStack contextStack = ThreadContext.cloneStack();
        
        return CompletableFuture.supplyAsync(() -> {
            // Restore context in async thread
            try (CloseableThreadContext.Instance ctc = 
                 CloseableThreadContext.putAll(contextMap)) {
                
                // Restore stack manually if needed
                for (String stackItem : contextStack.asList()) {
                    ThreadContext.push(stackItem);
                }
                
                try {
                    logger.info("Async operation with context");
                    return performAsyncOperation();
                } finally {
                    ThreadContext.clearStack();
                }
            }
        });
    }
}

// Executor service with context propagation
public class ContextAwareExecutor {
    private final ExecutorService executor = Executors.newFixedThreadPool(10);
    private static final Logger logger = LogManager.getLogger();
    
    public <T> CompletableFuture<T> executeWithContext(Supplier<T> task) {
        // Capture current thread context
        Map<String, String> contextMap = ThreadContext.getContext();
        ContextStack contextStack = ThreadContext.cloneStack();
        
        return CompletableFuture.supplyAsync(() -> {
            try (CloseableThreadContext.Instance ctc = 
                 CloseableThreadContext.putAll(contextMap)) {
                
                // Restore stack
                contextStack.asList().forEach(ThreadContext::push);
                
                try {
                    return task.get();
                } finally {
                    ThreadContext.clearStack();
                }
            }
        }, executor);
    }
}

// Spring async method with context
@Service
public class AsyncService {
    private static final Logger logger = LogManager.getLogger();
    
    @Async
    public CompletableFuture<String> processAsync(String data) {
        // Context must be manually propagated in @Async methods
        logger.info("Async processing started"); // May not have context
        
        try {
            String result = performProcessing(data);
            logger.info("Async processing completed");
            return CompletableFuture.completedFuture(result);
        } catch (Exception e) {
            logger.error("Async processing failed", e);
            throw e;
        }
    }
}