[#1658] Let trafficshaping work with ByteBuf and ByteBufHolder out of the box

This commit is contained in:
Norman Maurer 2013-07-29 07:31:09 +02:00
parent 6ce8571df3
commit 5e703e7419

View File

@ -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;
}
}