Fix regression in DefaultAttributeMap / Add Attribute.key() / More fine-grained locking

common/src/main/java/io/netty/util/Attribute.java

@@ -16,11 +16,17 @@
 package io.netty.util;
 
 /**
- * An attribute which allows to store an value reference. It may be updated atomically and so is thread-safe.
+ * An attribute which allows to store a value reference. It may be updated atomically and so is thread-safe.
  *
  * @param <T>   the type of the value it holds.
  */
 public interface Attribute<T> {
+
+    /**
+     * Returns the key of this attribute.
+     */
+    AttributeKey<T> key();
+
     /**
      * Returns the current value, which may be {@code null}
      */
@@ -42,6 +48,12 @@ public interface Attribute<T> {
      */
     T setIfAbsent(T value);
 
+    /**
+     * Removes this attribute from the {@link AttributeMap} and returns the old value..  Subsequent {@link #get()}
+     * calls will return @{code null}.
+     */
+    T getAndRemove();
+
     /**
      * Atomically sets the value to the given updated value if the current value == the expected value.
      * If it the set was successful it returns {@code true} otherwise {@code false}.
@@ -49,7 +61,7 @@ public interface Attribute<T> {
     boolean compareAndSet(T oldValue, T newValue);
 
     /**
-     * Remove this attribute from the {@link AttributeMap}.
+     * Removes this attribute from the {@link AttributeMap}.  Subsequent {@link #get()} calls will return @{code null}.
      */
     void remove();
 }

common/src/main/java/io/netty/util/DefaultAttributeMap.java

@@ -18,6 +18,7 @@ package io.netty.util;
 import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 
 /**
  * Default {@link AttributeMap} implementation which use simple synchronization to keep the memory overhead
@@ -25,30 +26,53 @@ import java.util.concurrent.atomic.AtomicReference;
  */
 public class DefaultAttributeMap implements AttributeMap {
 
-    // Initialize lazily to reduce memory consumption.
-    private Map<AttributeKey<?>, Attribute<?>> map;
+    @SuppressWarnings("rawtypes")
+    private static final AtomicReferenceFieldUpdater<DefaultAttributeMap, Map> updater =
+            AtomicReferenceFieldUpdater.newUpdater(DefaultAttributeMap.class, Map.class, "map");
+
+    // Initialize lazily to reduce memory consumption; updated by AtomicReferenceFieldUpdater above.
+    @SuppressWarnings("UnusedDeclaration")
+    private volatile Map<AttributeKey<?>, Attribute<?>> map;
 
     @Override
-    public synchronized <T> Attribute<T> attr(AttributeKey<T> key) {
+    public <T> Attribute<T> attr(AttributeKey<T> key) {
         Map<AttributeKey<?>, Attribute<?>> map = this.map;
         if (map == null) {
             // Not using ConcurrentHashMap due to high memory consumption.
-            map = this.map = new IdentityHashMap<AttributeKey<?>, Attribute<?>>(2);
+            map = new IdentityHashMap<AttributeKey<?>, Attribute<?>>(2);
+            if (!updater.compareAndSet(this, null, map)) {
+                map = this.map;
+            }
         }
 
-        @SuppressWarnings("unchecked")
-        Attribute<T> attr = (Attribute<T>) map.get(key);
-        if (attr == null) {
-            attr = new DefaultAttribute<T>();
-            map.put(key, attr);
+        synchronized (map) {
+            @SuppressWarnings("unchecked")
+            Attribute<T> attr = (Attribute<T>) map.get(key);
+            if (attr == null) {
+                attr = new DefaultAttribute<T>(map, key);
+                map.put(key, attr);
+            }
+            return attr;
         }
-        return attr;
     }
 
-    private final class DefaultAttribute<T> extends AtomicReference<T> implements Attribute<T> {
+    private static final class DefaultAttribute<T> extends AtomicReference<T> implements Attribute<T> {
 
         private static final long serialVersionUID = -2661411462200283011L;
 
+        private final Map<AttributeKey<?>, Attribute<?>> map;
+        private final AttributeKey<T> key;
+
+        DefaultAttribute(Map<AttributeKey<?>, Attribute<?>> map, AttributeKey<T> key) {
+            this.map = map;
+            this.key = key;
+        }
+
+        @Override
+        public AttributeKey<T> key() {
+            return key;
+        }
+
         @Override
         public T setIfAbsent(T value) {
             while (!compareAndSet(null, value)) {
@@ -60,10 +84,22 @@ public class DefaultAttributeMap implements AttributeMap {
             return null;
         }
 
+        @Override
+        public T getAndRemove() {
+            T oldValue = getAndSet(null);
+            remove0();
+            return oldValue;
+        }
+
         @Override
         public void remove() {
-            synchronized (DefaultAttributeMap.this) {
-                map.remove(this);
+            set(null);
+            remove0();
+        }
+
+        private void remove0() {
+            synchronized (map) {
+                map.remove(key);
             }
         }
     }