Expunge stale entries in InternalLoggerRegistry

Use ReferenceQueue to track and remove GC'd Logger WeakReferences from the registry map to prevent potential memory leaks. Includes improved unit tests.

Author: jhl221123

log4j-core-test/src/main/java/org/apache/logging/log4j/core/test/util/InternalLoggerRegistryGCTest.java

@@ -16,6 +16,9 @@
  */
 package org.apache.logging.log4j.core.test.util;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.awaitility.Awaitility.await;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -24,6 +27,8 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.lang.ref.WeakReference;
+import java.lang.reflect.Field;
+import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -32,17 +37,30 @@
 import org.apache.logging.log4j.core.util.internal.InternalLoggerRegistry;
 import org.apache.logging.log4j.message.MessageFactory;
 import org.apache.logging.log4j.message.SimpleMessageFactory;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInfo;
 
 class InternalLoggerRegistryTest {
+    private LoggerContext loggerContext;
     private InternalLoggerRegistry registry;
     private MessageFactory messageFactory;
 
     @BeforeEach
-    void setUp() {
-        registry = new InternalLoggerRegistry();
-        messageFactory = new SimpleMessageFactory();
+    void setUp(TestInfo testInfo) throws NoSuchFieldException, IllegalAccessException {
+        loggerContext = new LoggerContext(testInfo.getDisplayName());
+        Field registryField = loggerContext.getClass().getDeclaredField("loggerRegistry");
+        registryField.setAccessible(true);
+        registry = (InternalLoggerRegistry) registryField.get(loggerContext);
+        messageFactory = SimpleMessageFactory.INSTANCE;
+    }
+
+    @AfterEach
+    void tearDown() {
+        if (loggerContext != null) {
+            loggerContext.stop();
+        }
     }
 
     @Test
@@ -52,54 +70,79 @@ void testGetLoggerReturnsNullForNonExistentLogger() {
 
     @Test
     void testComputeIfAbsentCreatesLogger() {
-        Logger logger =
-                registry.computeIfAbsent("testLogger", messageFactory, (name, factory) -> LoggerContext.getContext()
-                        .getLogger(name, factory));
+        Logger logger = registry.computeIfAbsent(
+                "testLogger", messageFactory, (name, factory) -> loggerContext.getLogger(name, factory));
         assertNotNull(logger);
         assertEquals("testLogger", logger.getName());
     }
 
     @Test
     void testGetLoggerRetrievesExistingLogger() {
-        Logger logger =
-                registry.computeIfAbsent("testLogger", messageFactory, (name, factory) -> LoggerContext.getContext()
-                        .getLogger(name, factory));
+        Logger logger = registry.computeIfAbsent(
+                "testLogger", messageFactory, (name, factory) -> loggerContext.getLogger(name, factory));
         assertSame(logger, registry.getLogger("testLogger", messageFactory));
     }
 
     @Test
     void testHasLoggerReturnsCorrectStatus() {
         assertFalse(registry.hasLogger("testLogger", messageFactory));
-        registry.computeIfAbsent("testLogger", messageFactory, (name, factory) -> LoggerContext.getContext()
-                .getLogger(name, factory));
+        registry.computeIfAbsent(
+                "testLogger", messageFactory, (name, factory) -> loggerContext.getLogger(name, factory));
         assertTrue(registry.hasLogger("testLogger", messageFactory));
     }
 
     @Test
-    void testExpungeStaleEntriesRemovesGarbageCollectedLoggers() throws InterruptedException {
-        Logger logger =
-                registry.computeIfAbsent("testLogger", messageFactory, (name, factory) -> LoggerContext.getContext()
-                        .getLogger(name, factory));
-
-        WeakReference<Logger> weakRef = new WeakReference<>(logger);
-        logger = null; // Dereference to allow GC
+    void testExpungeStaleWeakReferenceEntries() {
+        String loggerNamePrefix = "testLogger_";
+        int numberOfLoggers = 1000;
+
+        for (int i = 0; i < numberOfLoggers; i++) {
+            Logger logger = registry.computeIfAbsent(
+                    loggerNamePrefix + i, messageFactory, (name, factory) -> loggerContext.getLogger(name, factory));
+            logger.info("Using logger {}", logger.getName());
+        }
 
-        // Retry loop to give GC time to collect
-        for (int i = 0; i < 10; i++) {
+        await().atMost(10, SECONDS).pollInterval(100, MILLISECONDS).untilAsserted(() -> {
             System.gc();
-            Thread.sleep(100);
-            if (weakRef.get() == null) {
-                break;
+            System.runFinalization();
+            registry.computeIfAbsent(
+                    "triggerExpunge", messageFactory, (name, factory) -> loggerContext.getLogger(name, factory));
+
+            Map<MessageFactory, Map<String, WeakReference<Logger>>> loggerRefByNameByMessageFactory =
+                    reflectAndGetLoggerMapFromRegistry();
+            Map<String, WeakReference<Logger>> loggerRefByName = loggerRefByNameByMessageFactory.get(messageFactory);
+
+            boolean isExpungeStaleEntries = true;
+            for (int i = 0; i < numberOfLoggers; i++) {
+                if (loggerRefByName.containsKey(loggerNamePrefix + i)) {
+                    isExpungeStaleEntries = false;
+                    break;
+                }
             }
-        }
+            assertTrue(
+                    isExpungeStaleEntries,
+                    "Stale WeakReference entries were not removed from the inner map for MessageFactory");
+        });
+    }
 
-        // Access the registry to potentially trigger cleanup
-        registry.computeIfAbsent("tempLogger", messageFactory, (name, factory) -> LoggerContext.getContext()
-                .getLogger(name, factory));
+    @Test
+    void testExpungeStaleMessageFactoryEntry() {
+        Logger logger = registry.computeIfAbsent(
+                "testLogger", messageFactory, (name, factory) -> loggerContext.getLogger(name, factory));
+        logger.info("Using logger {}", logger.getName());
+        logger = null;
 
-        assertNull(weakRef.get(), "Logger should have been garbage collected");
-        assertNull(
-                registry.getLogger("testLogger", messageFactory), "Stale logger should be removed from the registry");
+        await().atMost(10, SECONDS).pollInterval(100, MILLISECONDS).untilAsserted(() -> {
+            System.gc();
+            System.runFinalization();
+            registry.getLogger("testLogger", messageFactory);
+
+            Map<MessageFactory, Map<String, WeakReference<Logger>>> loggerRefByNameByMessageFactory =
+                    reflectAndGetLoggerMapFromRegistry();
+            assertNull(
+                    loggerRefByNameByMessageFactory.get(messageFactory),
+                    "Stale MessageFactory entry was not removed from the outer map");
+        });
     }
 
     @Test
@@ -110,8 +153,8 @@ void testConcurrentAccess() throws InterruptedException {
 
         for (int i = 0; i < threadCount; i++) {
             executor.submit(() -> {
-                registry.computeIfAbsent("testLogger", messageFactory, (name, factory) -> LoggerContext.getContext()
-                        .getLogger(name, factory));
+                registry.computeIfAbsent(
+                        "testLogger", messageFactory, (name, factory) -> loggerContext.getLogger(name, factory));
                 latch.countDown();
             });
         }
@@ -124,4 +167,14 @@ void testConcurrentAccess() throws InterruptedException {
                 registry.getLogger("testLogger", messageFactory),
                 "Logger should be accessible after concurrent creation");
     }
+
+    private Map<MessageFactory, Map<String, WeakReference<Logger>>> reflectAndGetLoggerMapFromRegistry()
+            throws NoSuchFieldException, IllegalAccessException {
+        Field loggerMapField = registry.getClass().getDeclaredField("loggerRefByNameByMessageFactory");
+        loggerMapField.setAccessible(true);
+        @SuppressWarnings("unchecked")
+        Map<MessageFactory, Map<String, WeakReference<Logger>>> loggerMap =
+                (Map<MessageFactory, Map<String, WeakReference<Logger>>>) loggerMapField.get(registry);
+        return loggerMap;
+    }
 }

log4j-core/src/main/java/org/apache/logging/log4j/core/util/internal/InternalLoggerRegistry.java

@@ -23,6 +23,7 @@
 import java.lang.ref.WeakReference;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.WeakHashMap;
 import java.util.concurrent.locks.Lock;
@@ -66,36 +67,31 @@ public final class InternalLoggerRegistry {
     public InternalLoggerRegistry() {}
 
     /**
-     * Expunges stale logger references from the registry.
+     * Expunges stale entries for logger references and message factories.
      */
     private void expungeStaleEntries() {
-        Reference<? extends Logger> loggerRef;
-        while ((loggerRef = staleLoggerRefs.poll()) != null) {
-            removeLogger(loggerRef);
-        }
-    }
-
-    /**
-     * Removes a logger from the registry.
-     */
-    private void removeLogger(Reference<? extends Logger> loggerRef) {
-        Logger logger = loggerRef.get();
-        if (logger == null) return; // Logger already cleared
+        Reference<? extends Logger> loggerRef = staleLoggerRefs.poll();
 
-        MessageFactory messageFactory = logger.getMessageFactory();
-        String name = logger.getName();
+        if (loggerRef != null) {
+            writeLock.lock();
+            try {
+                while (staleLoggerRefs.poll() != null) {
+                    // Clear refQueue
+                }
 
-        writeLock.lock();
-        try {
-            Map<String, WeakReference<Logger>> loggerRefByName = loggerRefByNameByMessageFactory.get(messageFactory);
-            if (loggerRefByName != null) {
-                loggerRefByName.remove(name);
-                if (loggerRefByName.isEmpty()) {
-                    loggerRefByNameByMessageFactory.remove(messageFactory); // Cleanup
+                Iterator<Map.Entry<MessageFactory, Map<String, WeakReference<Logger>>>> outerIt =
+                        loggerRefByNameByMessageFactory.entrySet().iterator();
+                while (outerIt.hasNext()) {
+                    Map.Entry<MessageFactory, Map<String, WeakReference<Logger>>> outerEntry = outerIt.next();
+                    Map<String, WeakReference<Logger>> innerMap = outerEntry.getValue();
+                    innerMap.values().removeIf(weakRef -> weakRef.get() == null);
+                    if (innerMap.isEmpty()) {
+                        outerIt.remove();
+                    }
                 }
+            } finally {
+                writeLock.unlock();
             }
-        } finally {
-            writeLock.unlock();
         }
     }