add a Netty ByteBuf based outputstream with consumer
This commit is contained in:
parent
41cca5c606
commit
fa3fa2e958
1 changed files with 132 additions and 0 deletions
|
@ -0,0 +1,132 @@
|
||||||
|
package org.xbib.net.http.server.netty.buffer;
|
||||||
|
|
||||||
|
import io.netty.buffer.ByteBuf;
|
||||||
|
import io.netty.buffer.ByteBufAllocator;
|
||||||
|
|
||||||
|
import java.io.IOException;
|
||||||
|
import java.io.OutputStream;
|
||||||
|
import java.util.concurrent.atomic.AtomicBoolean;
|
||||||
|
import java.util.function.Consumer;
|
||||||
|
|
||||||
|
public class NettyDataBufferOutputStream extends OutputStream {
|
||||||
|
|
||||||
|
private ByteBuf pooledBuffer;
|
||||||
|
|
||||||
|
private long written;
|
||||||
|
|
||||||
|
private final long maxContentLength;
|
||||||
|
|
||||||
|
private final AtomicBoolean closed;
|
||||||
|
|
||||||
|
private final AtomicBoolean writeStarted;
|
||||||
|
|
||||||
|
private final ByteBufAllocator byteBufAllocator;
|
||||||
|
|
||||||
|
private final Consumer<ByteBuf> consumer;
|
||||||
|
|
||||||
|
public NettyDataBufferOutputStream(ByteBufAllocator byteBufAllocator,
|
||||||
|
int maxContentLength,
|
||||||
|
Consumer<ByteBuf> consumer) {
|
||||||
|
this.byteBufAllocator = byteBufAllocator;
|
||||||
|
this.maxContentLength = maxContentLength;
|
||||||
|
this.consumer = consumer;
|
||||||
|
this.closed = new AtomicBoolean();
|
||||||
|
this.writeStarted = new AtomicBoolean();
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* {@inheritDoc}
|
||||||
|
*/
|
||||||
|
@Override
|
||||||
|
public void write(final int b) throws IOException {
|
||||||
|
write(new byte[]{(byte) b}, 0, 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* {@inheritDoc}
|
||||||
|
*/
|
||||||
|
@Override
|
||||||
|
public void write(final byte[] b) throws IOException {
|
||||||
|
write(b, 0, b.length);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* {@inheritDoc}
|
||||||
|
*/
|
||||||
|
@Override
|
||||||
|
public void write(final byte[] b, final int off, final int len) throws IOException {
|
||||||
|
if (len < 1) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if (closed.get()) {
|
||||||
|
throw new IOException("Stream is closed");
|
||||||
|
}
|
||||||
|
int remaining = len;
|
||||||
|
int index = off;
|
||||||
|
ByteBuf buffer = pooledBuffer;
|
||||||
|
try {
|
||||||
|
if (buffer == null) {
|
||||||
|
pooledBuffer = buffer = byteBufAllocator.buffer();
|
||||||
|
}
|
||||||
|
while (remaining > 0) {
|
||||||
|
int toWrite = Math.min(remaining, buffer.writableBytes());
|
||||||
|
buffer.writeBytes(b, index, toWrite);
|
||||||
|
remaining -= toWrite;
|
||||||
|
index += toWrite;
|
||||||
|
if (!buffer.isWritable()) {
|
||||||
|
writeStarted.set(true);
|
||||||
|
ByteBuf currentBuffer = buffer;
|
||||||
|
pooledBuffer = buffer = byteBufAllocator.buffer();
|
||||||
|
consumer.accept(currentBuffer);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} catch (Exception e) {
|
||||||
|
if (buffer != null) {
|
||||||
|
buffer.release();
|
||||||
|
}
|
||||||
|
throw new IOException(e);
|
||||||
|
}
|
||||||
|
this.written += len;
|
||||||
|
if (maxContentLength != -1 && this.written >= maxContentLength) {
|
||||||
|
close();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* {@inheritDoc}
|
||||||
|
*/
|
||||||
|
@Override
|
||||||
|
public void flush() throws IOException {
|
||||||
|
if (closed.get()) {
|
||||||
|
throw new IOException("Stream is closed");
|
||||||
|
}
|
||||||
|
internalFlush();
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* {@inheritDoc}
|
||||||
|
*/
|
||||||
|
@Override
|
||||||
|
public void close() throws IOException {
|
||||||
|
if (closed.get()) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
closed.set(true);
|
||||||
|
internalFlush();
|
||||||
|
}
|
||||||
|
|
||||||
|
private void internalFlush() throws IOException {
|
||||||
|
try {
|
||||||
|
if (pooledBuffer != null) {
|
||||||
|
consumer.accept(pooledBuffer);
|
||||||
|
}
|
||||||
|
} catch (Exception e) {
|
||||||
|
throw new IOException(e);
|
||||||
|
} finally {
|
||||||
|
if (pooledBuffer != null) {
|
||||||
|
pooledBuffer.release();
|
||||||
|
pooledBuffer = null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in a new issue