diff --git a/src/main/java/org/jboss/netty/channel/ChannelLocal.java b/src/main/java/org/jboss/netty/channel/ChannelLocal.java index 32abb6ae6a..72342323f3 100644 --- a/src/main/java/org/jboss/netty/channel/ChannelLocal.java +++ b/src/main/java/org/jboss/netty/channel/ChannelLocal.java @@ -40,13 +40,31 @@ public class ChannelLocal { private final ConcurrentMap map = new ConcurrentIdentityWeakKeyHashMap(); + + private final ChannelFutureListener remover = new ChannelFutureListener() { + @Override + public void operationComplete(ChannelFuture future) throws Exception { + remove(future.getChannel()); + } + }; + + private final boolean removeOnClose; + /** - * Creates a {@link Channel} local variable. + * Creates a {@link Channel} local variable by calling {@link #ChannelLocal(boolean)} with false as parameter */ public ChannelLocal() { - super(); + this(false); } + /** + * Creates a {@link Channel} local variable. + * + * @param removeOnClose if true the {@link ChannelLocal} will remove a {@link Channel} from it own once the {@link Channel} was closed. + */ + public ChannelLocal(boolean removeOnClose) { + this.removeOnClose = removeOnClose; + } /** * Returns the initial value of the variable. By default, it returns * {@code null}. Override it to change the initial value. @@ -88,7 +106,11 @@ public class ChannelLocal { if (channel == null) { throw new NullPointerException("channel"); } - return map.put(channel, value); + T old = map.put(channel, value); + if (removeOnClose) { + channel.getCloseFuture().addListener(remover); + } + return old; } } @@ -105,7 +127,12 @@ public class ChannelLocal { if (channel == null) { throw new NullPointerException("channel"); } - return map.putIfAbsent(channel, value); + T mapping = map.putIfAbsent(channel, value); + + if (removeOnClose && mapping == null) { + channel.getCloseFuture().addListener(remover); + } + return mapping; } } @@ -126,6 +153,9 @@ public class ChannelLocal { if (removed == null) { return initialValue(channel); } else { + if (removeOnClose) { + channel.getCloseFuture().removeListener(remover); + } return removed; } }