Use a map for the attachments.

Signed-off-by: James R. Perkins <[email protected]>
jboss-logging · Jun 27, 2023 · 9d6219b · 9d6219b
1 parent 11bea6f
commit 9d6219b
Show file tree

Hide file tree

Showing 2 changed files with 96 additions and 169 deletions.
diff --git a/src/main/java/org/jboss/logmanager/LogContext.java b/src/main/java/org/jboss/logmanager/LogContext.java
@@ -19,8 +19,9 @@
 
 package org.jboss.logmanager;
 
-import static org.jboss.logmanager.LoggerNode.attachmentsFull;
-
+import java.lang.invoke.ConstantBootstraps;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
 import java.security.AccessController;
 import java.security.Permission;
 import java.security.PrivilegedAction;
@@ -75,25 +76,10 @@ private static LogContextInitializer discoverDefaultInitializer0() {
     private final boolean strong;
     private final LogContextInitializer initializer;
 
-    /**
-     * The first attachment key.
-     */
-    private Logger.AttachmentKey<?> attachmentKey1;
+    private volatile Map<Logger.AttachmentKey<?>, Object> attachments;
 
-    /**
-     * The first attachment value.
-     */
-    private Object attachmentValue1;
-
-    /**
-     * The second attachment key.
-     */
-    private Logger.AttachmentKey<?> attachmentKey2;
-
-    /**
-     * The second attachment value.
-     */
-    private Object attachmentValue2;
+    private static final VarHandle attachmentHandle = ConstantBootstraps.fieldVarHandle(MethodHandles.lookup(), "attachments",
+            VarHandle.class, LogContext.class, Map.class);
 
     /**
      * This lazy holder class is required to prevent a problem due to a LogContext instance being constructed
@@ -147,6 +133,7 @@ private static void addStrong(Map<String, Reference<Level, Void>> map, Level lev
         levelMapReference = new AtomicReference<Map<String, Reference<Level, Void>>>(LazyHolder.INITIAL_LEVEL_MAP);
         rootLogger = new LoggerNode(this);
         closeHandlers = new LinkedHashSet<>();
+        attachments = Map.of();
     }
 
     /**
@@ -213,13 +200,7 @@ public static LogContext create(LogContextInitializer initializer) {
     @SuppressWarnings("unchecked")
     public <V> V getAttachment(Logger.AttachmentKey<V> key) {
         Assert.checkNotNullParam("key", key);
-        synchronized (this) {
-            if (key == attachmentKey1)
-                return (V) attachmentValue1;
-            if (key == attachmentKey2)
-                return (V) attachmentValue2;
-        }
-        return null;
+        return (V) attachments.get(key);
     }
 
     /**
@@ -240,26 +221,14 @@ public <V> V attach(Logger.AttachmentKey<V> key, V value) throws SecurityExcepti
         checkAccess();
         Assert.checkNotNullParam("key", key);
         Assert.checkNotNullParam("value", value);
+        Map<Logger.AttachmentKey<?>, Object> oldAttachments;
+        Map<Logger.AttachmentKey<?>, Object> newAttachments;
         V old;
-        synchronized (this) {
-            if (key == attachmentKey1) {
-                old = (V) attachmentValue1;
-                attachmentValue1 = value;
-            } else if (key == attachmentKey2) {
-                old = (V) attachmentValue2;
-                attachmentValue2 = value;
-            } else if (attachmentKey1 == null) {
-                old = null;
-                attachmentKey1 = key;
-                attachmentValue1 = value;
-            } else if (attachmentKey2 == null) {
-                old = null;
-                attachmentKey2 = key;
-                attachmentValue2 = value;
-            } else {
-                throw attachmentsFull();
-            }
-        }
+        do {
+            oldAttachments = attachments;
+            newAttachments = new HashMap<>(oldAttachments);
+            old = (V) newAttachments.put(key, value);
+        } while (!attachmentHandle.compareAndSet(this, oldAttachments, Map.copyOf(newAttachments)));
         return old;
     }
 
@@ -281,25 +250,17 @@ public <V> V attachIfAbsent(Logger.AttachmentKey<V> key, V value) throws Securit
         checkAccess();
         Assert.checkNotNullParam("key", key);
         Assert.checkNotNullParam("value", value);
-        V old;
-        synchronized (this) {
-            if (key == attachmentKey1) {
-                old = (V) attachmentValue1;
-            } else if (key == attachmentKey2) {
-                old = (V) attachmentValue2;
-            } else if (attachmentKey1 == null) {
-                old = null;
-                attachmentKey1 = key;
-                attachmentValue1 = value;
-            } else if (attachmentKey2 == null) {
-                old = null;
-                attachmentKey2 = key;
-                attachmentValue2 = value;
-            } else {
-                throw attachmentsFull();
+        Map<Logger.AttachmentKey<?>, Object> oldAttachments;
+        Map<Logger.AttachmentKey<?>, Object> newAttachments;
+        do {
+            oldAttachments = attachments;
+            if (oldAttachments.containsKey(key)) {
+                return (V) oldAttachments.get(key);
             }
-        }
-        return old;
+            newAttachments = new HashMap<>(oldAttachments);
+            newAttachments.put(key, value);
+        } while (!attachmentHandle.compareAndSet(this, oldAttachments, Map.copyOf(newAttachments)));
+        return null;
     }
 
     /**
@@ -315,19 +276,25 @@ public <V> V attachIfAbsent(Logger.AttachmentKey<V> key, V value) throws Securit
     public <V> V detach(Logger.AttachmentKey<V> key) throws SecurityException {
         checkAccess();
         Assert.checkNotNullParam("key", key);
-        V old;
-        synchronized (this) {
-            if (key == attachmentKey1) {
-                old = (V) attachmentValue1;
-                attachmentValue1 = null;
-            } else if (key == attachmentKey2) {
-                old = (V) attachmentValue2;
-                attachmentValue2 = null;
+        Map<Logger.AttachmentKey<?>, Object> oldAttachments;
+        Map<Logger.AttachmentKey<?>, Object> newAttachments;
+        V result;
+        do {
+            oldAttachments = attachments;
+            result = (V) oldAttachments.get(key);
+            if (result == null) {
+                return null;
+            }
+            final int size = oldAttachments.size();
+            if (size == 1) {
+                // special case - the new map is empty
+                newAttachments = Map.of();
             } else {
-                old = null;
+                newAttachments = new HashMap<>(oldAttachments);
+                newAttachments.remove(key);
             }
-        }
-        return old;
+        } while (!attachmentHandle.compareAndSet(this, oldAttachments, Map.copyOf(newAttachments)));
+        return result;
     }
 
     /**
@@ -529,18 +496,14 @@ public void close() throws Exception {
             for (AutoCloseable handler : closeHandlers) {
                 handler.close();
             }
-            synchronized (this) {
-                attachmentKey1 = null;
-                attachmentKey2 = null;
-                final var value1 = attachmentValue1;
-                attachmentValue1 = null;
-                if (value1 instanceof AutoCloseable) {
-                    ((AutoCloseable) value1).close();
-                }
-                final var value2 = attachmentValue2;
-                attachmentValue2 = null;
-                if (value2 instanceof AutoCloseable) {
-                    ((AutoCloseable) value2).close();
+            @SuppressWarnings("unchecked")
+            final Map<Object, Object> oldAttachments = (Map<Object, Object>) attachmentHandle.getAndSet(this, Map.of());
+            for (Object value : oldAttachments.values()) {
+                if (value instanceof AutoCloseable) {
+                    try {
+                        ((AutoCloseable) value).close();
+                    } catch (Exception ignore) {
+                    }
                 }
             }
         } finally {