Examples with PooledByteBuffer io.undertow.connector.PooledByteBuffer used on opensource projects

Example 31 with PooledByteBuffer

use of io.undertow.connector.PooledByteBuffer in project undertow by undertow-io.

the class Transfer method initiateTransfer.

/**
     * Initiate a low-copy transfer between two stream channels.  The pool should be a direct buffer pool for best
     * performance.
     *
     * @param source the source channel
     * @param sink the target channel
     * @param sourceListener the source listener to set and call when the transfer is complete, or {@code null} to clear the listener at that time
     * @param sinkListener the target listener to set and call when the transfer is complete, or {@code null} to clear the listener at that time
     * @param readExceptionHandler the read exception handler to call if an error occurs during a read operation
     * @param writeExceptionHandler the write exception handler to call if an error occurs during a write operation
     * @param pool the pool from which the transfer buffer should be allocated
     */
public static <I extends StreamSourceChannel, O extends StreamSinkChannel> void initiateTransfer(final I source, final O sink, final ChannelListener<? super I> sourceListener, final ChannelListener<? super O> sinkListener, final ChannelExceptionHandler<? super I> readExceptionHandler, final ChannelExceptionHandler<? super O> writeExceptionHandler, ByteBufferPool pool) {
    if (pool == null) {
        throw UndertowMessages.MESSAGES.argumentCannotBeNull("pool");
    }
    final PooledByteBuffer allocated = pool.allocate();
    boolean free = true;
    try {
        final ByteBuffer buffer = allocated.getBuffer();
        long read;
        for (; ; ) {
            try {
                read = source.read(buffer);
                buffer.flip();
            } catch (IOException e) {
                ChannelListeners.invokeChannelExceptionHandler(source, readExceptionHandler, e);
                return;
            }
            if (read == 0 && !buffer.hasRemaining()) {
                break;
            }
            if (read == -1 && !buffer.hasRemaining()) {
                done(source, sink, sourceListener, sinkListener);
                return;
            }
            while (buffer.hasRemaining()) {
                final int res;
                try {
                    res = sink.write(buffer);
                } catch (IOException e) {
                    ChannelListeners.invokeChannelExceptionHandler(sink, writeExceptionHandler, e);
                    return;
                }
                if (res == 0) {
                    break;
                }
            }
            if (buffer.hasRemaining()) {
                break;
            }
            buffer.clear();
        }
        PooledByteBuffer current = null;
        if (buffer.hasRemaining()) {
            current = allocated;
            free = false;
        }
        final TransferListener<I, O> listener = new TransferListener<I, O>(pool, current, source, sink, sourceListener, sinkListener, writeExceptionHandler, readExceptionHandler, read == -1);
        sink.getWriteSetter().set(listener);
        source.getReadSetter().set(listener);
        //we resume both reads and writes, as we want to keep trying to fill the buffer
        if (current == null || buffer.capacity() != buffer.remaining()) {
            //we don't resume if the buffer is 100% full
            source.resumeReads();
        }
        if (current != null) {
            //we don't resume writes if we have nothing to write
            sink.resumeWrites();
        }
    } finally {
        if (free) {
            allocated.close();
        }
    }
}

Example 32 with PooledByteBuffer

use of io.undertow.connector.PooledByteBuffer in project undertow by undertow-io.

the class PathResource method serveImpl.

private void serveImpl(final Sender sender, final HttpServerExchange exchange, final long start, final long end, final IoCallback callback, final boolean range) {
    abstract class BaseFileTask implements Runnable {

        protected volatile FileChannel fileChannel;

        protected boolean openFile() {
            try {
                fileChannel = FileChannel.open(file, StandardOpenOption.READ);
                if (range) {
                    fileChannel.position(start);
                }
            } catch (NoSuchFileException e) {
                exchange.setStatusCode(StatusCodes.NOT_FOUND);
                callback.onException(exchange, sender, e);
                return false;
            } catch (IOException e) {
                exchange.setStatusCode(StatusCodes.INTERNAL_SERVER_ERROR);
                callback.onException(exchange, sender, e);
                return false;
            }
            return true;
        }
    }
    class ServerTask extends BaseFileTask implements IoCallback {

        private PooledByteBuffer pooled;

        long remaining = end - start + 1;

        @Override
        public void run() {
            if (range && remaining == 0) {
                //we are done
                pooled.close();
                pooled = null;
                IoUtils.safeClose(fileChannel);
                callback.onComplete(exchange, sender);
                return;
            }
            if (fileChannel == null) {
                if (!openFile()) {
                    return;
                }
                pooled = exchange.getConnection().getByteBufferPool().allocate();
            }
            if (pooled != null) {
                ByteBuffer buffer = pooled.getBuffer();
                try {
                    buffer.clear();
                    int res = fileChannel.read(buffer);
                    if (res == -1) {
                        //we are done
                        pooled.close();
                        IoUtils.safeClose(fileChannel);
                        callback.onComplete(exchange, sender);
                        return;
                    }
                    buffer.flip();
                    if (range) {
                        if (buffer.remaining() > remaining) {
                            buffer.limit((int) (buffer.position() + remaining));
                        }
                        remaining -= buffer.remaining();
                    }
                    sender.send(buffer, this);
                } catch (IOException e) {
                    onException(exchange, sender, e);
                }
            }
        }

        @Override
        public void onComplete(final HttpServerExchange exchange, final Sender sender) {
            if (exchange.isInIoThread()) {
                exchange.dispatch(this);
            } else {
                run();
            }
        }

        @Override
        public void onException(final HttpServerExchange exchange, final Sender sender, final IOException exception) {
            UndertowLogger.REQUEST_IO_LOGGER.ioException(exception);
            if (pooled != null) {
                pooled.close();
                pooled = null;
            }
            IoUtils.safeClose(fileChannel);
            if (!exchange.isResponseStarted()) {
                exchange.setStatusCode(StatusCodes.INTERNAL_SERVER_ERROR);
            }
            callback.onException(exchange, sender, exception);
        }
    }
    class TransferTask extends BaseFileTask {

        @Override
        public void run() {
            if (!openFile()) {
                return;
            }
            sender.transferFrom(fileChannel, new IoCallback() {

                @Override
                public void onComplete(HttpServerExchange exchange, Sender sender) {
                    try {
                        IoUtils.safeClose(fileChannel);
                    } finally {
                        callback.onComplete(exchange, sender);
                    }
                }

                @Override
                public void onException(HttpServerExchange exchange, Sender sender, IOException exception) {
                    try {
                        IoUtils.safeClose(fileChannel);
                    } finally {
                        callback.onException(exchange, sender, exception);
                    }
                }
            });
        }
    }
    BaseFileTask task;
    try {
        task = manager.getTransferMinSize() > Files.size(file) || range ? new ServerTask() : new TransferTask();
    } catch (IOException e) {
        throw new RuntimeException(e);
    }
    if (exchange.isInIoThread()) {
        exchange.dispatch(task);
    } else {
        task.run();
    }
}

Example 33 with PooledByteBuffer

use of io.undertow.connector.PooledByteBuffer in project undertow by undertow-io.

the class ChunkedStreamSinkConduit method createLastChunk.

private void createLastChunk(final boolean writeFinal) throws UnsupportedEncodingException {
    PooledByteBuffer lastChunkBufferPooled = bufferPool.allocate();
    ByteBuffer lastChunkBuffer = lastChunkBufferPooled.getBuffer();
    if (writeFinal) {
        lastChunkBuffer.put(CRLF);
    } else if (chunkingSepBuffer.hasRemaining()) {
        //the end of chunk /r/n has not been written yet
        //just add it to this buffer to make managing state easier
        lastChunkBuffer.put(chunkingSepBuffer);
    }
    lastChunkBuffer.put(LAST_CHUNK);
    //we just assume it will fit
    HeaderMap trailers = attachable.getAttachment(HttpAttachments.RESPONSE_TRAILERS);
    if (trailers != null && trailers.size() != 0) {
        for (HeaderValues trailer : trailers) {
            for (String val : trailer) {
                trailer.getHeaderName().appendTo(lastChunkBuffer);
                lastChunkBuffer.put((byte) ':');
                lastChunkBuffer.put((byte) ' ');
                lastChunkBuffer.put(val.getBytes(StandardCharsets.US_ASCII));
                lastChunkBuffer.put(CRLF);
            }
        }
        lastChunkBuffer.put(CRLF);
    } else {
        lastChunkBuffer.put(CRLF);
    }
    //horrible hack
    //there is a situation where we can get a buffer leak here if the connection is terminated abnormaly
    //this should be fixed once this channel has its lifecycle tied to the connection, same as fixed length
    lastChunkBuffer.flip();
    ByteBuffer data = ByteBuffer.allocate(lastChunkBuffer.remaining());
    data.put(lastChunkBuffer);
    data.flip();
    this.lastChunkBuffer = new ImmediatePooledByteBuffer(data);
    lastChunkBufferPooled.close();
}

Example 34 with PooledByteBuffer

use of io.undertow.connector.PooledByteBuffer in project undertow by undertow-io.

the class ChunkedStreamSourceConduit method read.

@Override
public int read(final ByteBuffer dst) throws IOException {
    try {
        long chunkRemaining = chunkReader.getChunkRemaining();
        //we have read the last chunk, we just return EOF
        if (chunkRemaining == -1) {
            return -1;
        }
        if (closed) {
            throw new ClosedChannelException();
        }
        PooledByteBuffer pooled = bufferWrapper.allocate();
        ByteBuffer buf = pooled.getBuffer();
        boolean free = true;
        try {
            //we need to do our initial read into a
            int r = next.read(buf);
            buf.flip();
            if (r == -1) {
                //Channel is broken, not sure how best to report it
                throw new ClosedChannelException();
            } else if (r == 0) {
                return 0;
            }
            if (chunkRemaining == 0) {
                chunkRemaining = chunkReader.readChunk(buf);
                if (chunkRemaining <= 0) {
                    return (int) chunkRemaining;
                }
            }
            final int originalLimit = dst.limit();
            try {
                //now we may have some stuff in the raw buffer
                //or the raw buffer may be exhausted, and we should read directly into the destination buffer
                //from the next
                int read = 0;
                long chunkInBuffer = Math.min(buf.remaining(), chunkRemaining);
                int remaining = dst.remaining();
                if (chunkInBuffer > remaining) {
                    //it won't fit
                    int orig = buf.limit();
                    buf.limit(buf.position() + remaining);
                    dst.put(buf);
                    buf.limit(orig);
                    chunkRemaining -= remaining;
                    updateRemainingAllowed(remaining);
                    free = false;
                    return remaining;
                } else if (buf.hasRemaining()) {
                    int old = buf.limit();
                    buf.limit((int) Math.min(old, buf.position() + chunkInBuffer));
                    try {
                        dst.put(buf);
                    } finally {
                        buf.limit(old);
                    }
                    read += chunkInBuffer;
                    chunkRemaining -= chunkInBuffer;
                }
                //adjusting the limit as necessary to make sure we do not read too much
                if (chunkRemaining > 0) {
                    int old = dst.limit();
                    try {
                        if (chunkRemaining < dst.remaining()) {
                            dst.limit((int) (dst.position() + chunkRemaining));
                        }
                        int c = 0;
                        do {
                            c = next.read(dst);
                            if (c > 0) {
                                read += c;
                                chunkRemaining -= c;
                            }
                        } while (c > 0 && chunkRemaining > 0);
                        if (c == -1) {
                            throw new ClosedChannelException();
                        }
                    } finally {
                        dst.limit(old);
                    }
                } else {
                    free = false;
                }
                updateRemainingAllowed(read);
                return read;
            } finally {
                //buffer will be freed if not needed in exitRead
                dst.limit(originalLimit);
            }
        } finally {
            if (chunkRemaining >= 0) {
                chunkReader.setChunkRemaining(chunkRemaining);
            }
            if (!free && buf.hasRemaining()) {
                bufferWrapper.pushBack(pooled);
            } else {
                pooled.close();
            }
        }
    } catch (IOException | RuntimeException e) {
        IoUtils.safeClose(exchange.getConnection());
        throw e;
    }
}

Example 35 with PooledByteBuffer

use of io.undertow.connector.PooledByteBuffer in project undertow by undertow-io.

the class DeflatingStreamSinkConduit method deflateData.

/**
     * Runs the current data through the deflater. As much as possible this will be buffered in the current output
     * stream.
     *
     * @throws IOException
     */
private void deflateData(boolean force) throws IOException {
    //we don't need to flush here, as this should have been called already by the time we get to
    //this point
    boolean nextCreated = false;
    try (PooledByteBuffer arrayPooled = this.exchange.getConnection().getByteBufferPool().getArrayBackedPool().allocate()) {
        PooledByteBuffer pooled = this.currentBuffer;
        final ByteBuffer outputBuffer = pooled.getBuffer();
        final boolean shutdown = anyAreSet(state, SHUTDOWN);
        ByteBuffer buf = arrayPooled.getBuffer();
        while (force || !deflater.needsInput() || (shutdown && !deflater.finished())) {
            int count = deflater.deflate(buf.array(), buf.arrayOffset(), buf.remaining(), force ? Deflater.SYNC_FLUSH : Deflater.NO_FLUSH);
            Connectors.updateResponseBytesSent(exchange, count);
            if (count != 0) {
                int remaining = outputBuffer.remaining();
                if (remaining > count) {
                    outputBuffer.put(buf.array(), buf.arrayOffset(), count);
                } else {
                    if (remaining == count) {
                        outputBuffer.put(buf.array(), buf.arrayOffset(), count);
                    } else {
                        outputBuffer.put(buf.array(), buf.arrayOffset(), remaining);
                        additionalBuffer = ByteBuffer.allocate(count - remaining);
                        additionalBuffer.put(buf.array(), buf.arrayOffset() + remaining, count - remaining);
                        additionalBuffer.flip();
                    }
                    outputBuffer.flip();
                    this.state |= FLUSHING_BUFFER;
                    if (next == null) {
                        nextCreated = true;
                        this.next = createNextChannel();
                    }
                    if (!performFlushIfRequired()) {
                        return;
                    }
                }
            } else {
                force = false;
            }
        }
    } finally {
        if (nextCreated) {
            if (anyAreSet(state, WRITES_RESUMED)) {
                next.resumeWrites();
            }
        }
    }
}