<commit_before>package io.undertow.conduits;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.Channel;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.FileChannel;
import java.util.concurrent.TimeUnit;

import io.undertow.UndertowLogger;
import io.undertow.server.ConduitWrapper;
import io.undertow.server.ExchangeCompletionListener;
import io.undertow.server.HttpServerConnection;
import io.undertow.server.HttpServerExchange;
import io.undertow.util.AttachmentKey;
import io.undertow.util.ConduitFactory;
import org.xnio.Buffers;
import org.xnio.ChannelListener;
import org.xnio.IoUtils;
import org.xnio.Pool;
import org.xnio.Pooled;
import org.xnio.channels.ConnectedStreamChannel;
import org.xnio.channels.StreamSourceChannel;
import org.xnio.conduits.AbstractStreamSinkConduit;
import org.xnio.conduits.ConduitWritableByteChannel;
import org.xnio.conduits.StreamSinkConduit;

import static org.xnio.Bits.anyAreClear;
import static org.xnio.Bits.anyAreSet;

/**
 * A buffer that is used when processing pipelined requests, that allows the server to
 * buffer multiple responses into a single write() call.
 * <p/>
 * This can improve performance when pipelining requests.
 *
 * @author <NAME>
 */
public class PipelingBufferingStreamSinkConduit extends AbstractStreamSinkConduit<StreamSinkConduit> {

    public static final AttachmentKey<PipelingBufferingStreamSinkConduit> ATTACHMENT_KEY = AttachmentKey.create(PipelingBufferingStreamSinkConduit.class);

    /**
     * If this channel is shutdown
     */
    private static final int SHUTDOWN = 1;
    private static final int DELEGATE_SHUTDOWN = 1 << 1;
    private static final int FLUSHING = 1 << 3;

    private int state;

    private final Pool<ByteBuffer> pool;
    private Pooled<ByteBuffer> buffer;

    private final ExchangeCompletionListener completionListener = new ExchangeCompletionListener() {
        @Override
        public void exchangeEvent(final HttpServerExchange exchange, final NextListener nextListener) {
            //if we ever fail to read then we flush the pipeline buffer
            //this relies on us always doing an eager read when starting a request,
            //rather than waiting to be notified of data being available
            final HttpServerConnection connection = exchange.getConnection();
            if (connection.getExtraBytes() == null || exchange.isUpgrade()) {
                try {
                    if (!flushPipelinedData()) {
                        final ConnectedStreamChannel channel = connection.getChannel();
                        channel.getWriteSetter().set(new ChannelListener<Channel>() {
                            @Override
                            public void handleEvent(Channel c) {
                                try {
                                    if (flushPipelinedData()) {
                                        channel.getWriteSetter().set(null);
                                        channel.suspendWrites();
                                        nextListener.proceed();
                                    }
                                } catch (IOException e) {
                                    UndertowLogger.REQUEST_LOGGER.exceptionProcessingRequest(e);
                                    IoUtils.safeClose(channel);
                                }
                            }
                        });
                        connection.getChannel().resumeWrites();
                        return;
                    } else {
                        nextListener.proceed();
                    }
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            } else {
                nextListener.proceed();
            }
        }
    };

    public PipelingBufferingStreamSinkConduit(StreamSinkConduit next, final Pool<ByteBuffer> pool) {
        super(next);
        this.pool = pool;
    }

    @Override
    public long transferFrom(FileChannel src, long position, long count) throws IOException {
        if (anyAreSet(state, SHUTDOWN)) {
            throw new ClosedChannelException();
        }
        return src.transferTo(position, count, new ConduitWritableByteChannel(this));
    }

    @Override
    public long transferFrom(StreamSourceChannel source, long count, ByteBuffer throughBuffer) throws IOException {
        return IoUtils.transfer(source, count, throughBuffer, new ConduitWritableByteChannel(this));
    }

    @Override
    public long write(ByteBuffer[] srcs, int offset, int length) throws IOException {
        if (anyAreSet(state, SHUTDOWN)) {
            throw new ClosedChannelException();
        }
        if (anyAreSet(state, FLUSHING)) {
            boolean res = flushBuffer();
            if (!res) {
                return 0;
            }
        }
        Pooled<ByteBuffer> pooled = this.buffer;
        if (pooled == null) {
            this.buffer = pooled = pool.allocate();
        }
        final ByteBuffer buffer = pooled.getResource();

        int total = 0;
        for (int i = offset; i < offset + length; ++i) {
            total += srcs[i].remaining();
        }

        if (buffer.remaining() > total) {
            int put = total;
            Buffers.copy(buffer, srcs, offset, length);
            return put;
        } else {
            int put = buffer.remaining();
            Buffers.copy(put, buffer, srcs, offset, length);
            flushBuffer();
            return put;
        }
    }

    @Override
    public int write(ByteBuffer src) throws IOException {
        if (anyAreSet(state, SHUTDOWN)) {
            throw new ClosedChannelException();
        }
        if (anyAreSet(state, FLUSHING)) {
            boolean res = flushBuffer();
            if (!res) {
                return 0;
            }
        }
        Pooled<ByteBuffer> pooled = this.buffer;
        if (pooled == null) {
            this.buffer = pooled = pool.allocate();
        }
        final ByteBuffer buffer = pooled.getResource();
        if (buffer.remaining() > src.remaining()) {
            int put = src.remaining();
            buffer.put(src);
            return put;
        } else {
            int put = buffer.remaining();
            int old = src.limit();
            src.limit(src.position() + put);
            buffer.put(src);
            src.limit(old);
            flushBuffer();
            return put;
        }
    }

    /**
     * Flushes the cached data.
     * <p/>
     * This should be called when a read thread fails to read any more request data, to make sure that any
     * buffered data is flushed after the last pipelined request.
     * <p/>
     * If this returns false the read thread should suspend reads and resume writes
     *
     * @return <code>true</code> If the flush suceeded, false otherwise
     * @throws IOException
     */
    public boolean flushPipelinedData() throws IOException {
        if (buffer == null || buffer.getResource().position() == 0) {
            return next.flush();
        }
        if (!flushBuffer()) {
            return false;
        }
        return next.flush();
    }

    /**
     * Gets the channel wrapper that implements the buffering
     *
     * @return The channel wrapper
     */
    public ConduitWrapper<StreamSinkConduit> getChannelWrapper() {
        return new ConduitWrapper<StreamSinkConduit>() {
            @Override
            public StreamSinkConduit wrap(ConduitFactory<StreamSinkConduit> factory, HttpServerExchange exchange) {
                exchange.addExchangeCompleteListener(completionListener);
                return PipelingBufferingStreamSinkConduit.this;
            }<commit_msg>Flush the pipelining buffer correctly<commit_after>package io.undertow.conduits;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.Channel;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.FileChannel;
import java.util.concurrent.TimeUnit;

import io.undertow.UndertowLogger;
import io.undertow.server.ConduitWrapper;
import io.undertow.server.ExchangeCompletionListener;
import io.undertow.server.HttpServerConnection;
import io.undertow.server.HttpServerExchange;
import io.undertow.util.AttachmentKey;
import io.undertow.util.ConduitFactory;
import org.xnio.Buffers;
import org.xnio.ChannelListener;
import org.xnio.IoUtils;
import org.xnio.Pool;
import org.xnio.Pooled;
import org.xnio.channels.ConnectedStreamChannel;
import org.xnio.channels.StreamSourceChannel;
import org.xnio.conduits.AbstractStreamSinkConduit;
import org.xnio.conduits.ConduitWritableByteChannel;
import org.xnio.conduits.StreamSinkConduit;

import static org.xnio.Bits.allAreClear;
import static org.xnio.Bits.anyAreClear;
import static org.xnio.Bits.anyAreSet;

/**
 * A buffer that is used when processing pipelined requests, that allows the server to
 * buffer multiple responses into a single write() call.
 * <p/>
 * This can improve performance when pipelining requests.
 *
 * @author <NAME>
 */
public class PipelingBufferingStreamSinkConduit extends AbstractStreamSinkConduit<StreamSinkConduit> {

    public static final AttachmentKey<PipelingBufferingStreamSinkConduit> ATTACHMENT_KEY = AttachmentKey.create(PipelingBufferingStreamSinkConduit.class);

    /**
     * If this channel is shutdown
     */
    private static final int SHUTDOWN = 1;
    private static final int DELEGATE_SHUTDOWN = 1 << 1;
    private static final int FLUSHING = 1 << 3;

    private int state;

    private final Pool<ByteBuffer> pool;
    private Pooled<ByteBuffer> buffer;

    private final ExchangeCompletionListener completionListener = new ExchangeCompletionListener() {
        @Override
        public void exchangeEvent(final HttpServerExchange exchange, final NextListener nextListener) {
            //if we ever fail to read then we flush the pipeline buffer
            //this relies on us always doing an eager read when starting a request,
            //rather than waiting to be notified of data being available
            final HttpServerConnection connection = exchange.getConnection();
            if (connection.getExtraBytes() == null || exchange.isUpgrade()) {
                try {
                    if (!flushPipelinedData()) {
                        final ConnectedStreamChannel channel = connection.getChannel();
                        channel.getWriteSetter().set(new ChannelListener<Channel>() {
                            @Override
                            public void handleEvent(Channel c) {
                                try {
                                    if (flushPipelinedData()) {
                                        channel.getWriteSetter().set(null);
                                        channel.suspendWrites();
                                        nextListener.proceed();
                                    }
                                } catch (IOException e) {
                                    UndertowLogger.REQUEST_LOGGER.exceptionProcessingRequest(e);
                                    IoUtils.safeClose(channel);
                                }
                            }
                        });
                        connection.getChannel().resumeWrites();
                        return;
                    } else {
                        nextListener.proceed();
                    }
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            } else {
                nextListener.proceed();
            }
        }
    };

    public PipelingBufferingStreamSinkConduit(StreamSinkConduit next, final Pool<ByteBuffer> pool) {
        super(next);
        this.pool = pool;
    }

    @Override
    public long transferFrom(FileChannel src, long position, long count) throws IOException {
        if (anyAreSet(state, SHUTDOWN)) {
            throw new ClosedChannelException();
        }
        return src.transferTo(position, count, new ConduitWritableByteChannel(this));
    }

    @Override
    public long transferFrom(StreamSourceChannel source, long count, ByteBuffer throughBuffer) throws IOException {
        return IoUtils.transfer(source, count, throughBuffer, new ConduitWritableByteChannel(this));
    }

    @Override
    public long write(ByteBuffer[] srcs, int offset, int length) throws IOException {
        if (anyAreSet(state, SHUTDOWN)) {
            throw new ClosedChannelException();
        }
        if (anyAreSet(state, FLUSHING)) {
            boolean res = flushBuffer();
            if (!res) {
                return 0;
            }
        }
        Pooled<ByteBuffer> pooled = this.buffer;
        if (pooled == null) {
            this.buffer = pooled = pool.allocate();
        }
        final ByteBuffer buffer = pooled.getResource();

        int total = 0;
        for (int i = offset; i < offset + length; ++i) {
            total += srcs[i].remaining();
        }

        if (buffer.remaining() > total) {
            int put = total;
            Buffers.copy(buffer, srcs, offset, length);
            return put;
        } else {
            int put = buffer.remaining();
            Buffers.copy(put, buffer, srcs, offset, length);
            flushBuffer();
            return put;
        }
    }

    @Override
    public int write(ByteBuffer src) throws IOException {
        if (anyAreSet(state, SHUTDOWN)) {
            throw new ClosedChannelException();
        }
        if (anyAreSet(state, FLUSHING)) {
            boolean res = flushBuffer();
            if (!res) {
                return 0;
            }
        }
        Pooled<ByteBuffer> pooled = this.buffer;
        if (pooled == null) {
            this.buffer = pooled = pool.allocate();
        }
        final ByteBuffer buffer = pooled.getResource();
        if (buffer.remaining() > src.remaining()) {
            int put = src.remaining();
            buffer.put(src);
            return put;
        } else {
            int put = buffer.remaining();
            int old = src.limit();
            src.limit(src.position() + put);
            buffer.put(src);
            src.limit(old);
            flushBuffer();
            return put;
        }
    }

    /**
     * Flushes the cached data.
     * <p/>
     * This should be called when a read thread fails to read any more request data, to make sure that any
     * buffered data is flushed after the last pipelined request.
     * <p/>
     * If this returns false the read thread should suspend reads and resume writes
     *
     * @return <code>true</code> If the flush suceeded, false otherwise
     * @throws IOException
     */
    public boolean flushPipelinedData() throws IOException {
        if (buffer == null || (buffer.getResource().position() == 0 && allAreClear(state, FLUSHING))) {
            return next.flush();
        }
        return flushBuffer();
    }

    /**
     * Gets the channel wrapper that implements the buffering
     *
     * @return The channel wrapper
     */
    public ConduitWrapper<StreamSinkConduit> getChannelWrapper() {
        return new ConduitWrapper<StreamSinkConduit>() {
            @Override
            public StreamSinkConduit wrap(ConduitFactory<StreamSinkConduit> factory, HttpServerExchange exchange) {
                exchange.addExchangeCompleteListener(completionListener);
                return PipelingBufferingStreamSinkConduit.this;
            }
{'fraction_non_alphanumeric': 0.08590932509015971, 'fraction_numerical': 0.00115919629057187, 'mean_word_length': 1.9675138543856296, 'pattern_counts': {'":': 0, '<': 47, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6081132', 'n_tokens_mistral': 4006, 'n_tokens_neox': 3868, 'n_words': 1035}