diff --git a/handler/src/main/java/io/netty/handler/traffic/AbstractTrafficShapingHandler.java b/handler/src/main/java/io/netty/handler/traffic/AbstractTrafficShapingHandler.java index 05586e0a33..35343e80ef 100644 --- a/handler/src/main/java/io/netty/handler/traffic/AbstractTrafficShapingHandler.java +++ b/handler/src/main/java/io/netty/handler/traffic/AbstractTrafficShapingHandler.java @@ -16,6 +16,7 @@ package io.netty.handler.traffic; import io.netty.buffer.ByteBuf; +import io.netty.buffer.ByteBufHolder; import io.netty.channel.ChannelDuplexHandler; import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelPromise; @@ -279,9 +280,9 @@ public abstract class AbstractTrafficShapingHandler extends ChannelDuplexHandler public void write(final ChannelHandlerContext ctx, final Object msg, ChannelPromise promise) throws Exception { long curtime = System.currentTimeMillis(); - long size = ((ByteBuf) msg).readableBytes(); + long size = calculateSize(msg); - if (trafficCounter != null) { + if (size > -1 && trafficCounter != null) { trafficCounter.bytesWriteFlowControl(size); if (writeLimit == 0) { ctx.write(msg); @@ -327,4 +328,21 @@ public abstract class AbstractTrafficShapingHandler extends ChannelDuplexHandler " Read Limit: " + readLimit + " and Counter: " + (trafficCounter != null? trafficCounter.toString() : "none"); } + + /** + * Calculate the size of the given {@link Object}. + * + * This implementation supports {@link ByteBuf} and {@link ByteBufHolder}. Sub-classes may override this. + * @param msg the msg for which the size should be calculated + * @return size the size of the msg or {@code -1} if unknown. + */ + protected long calculateSize(Object msg) { + if (msg instanceof ByteBuf) { + return ((ByteBuf) msg).readableBytes(); + } + if (msg instanceof ByteBufHolder) { + return ((ByteBufHolder) msg).content().readableBytes(); + } + return -1; + } }