src/java.net.http/share/classes/jdk/internal/net/http/BufferingSubscriber.java

/*
 * Copyright (c) 2017, 2023, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.  Oracle designates this
 * particular file as subject to the "Classpath" exception as provided
 * by Oracle in the LICENSE file that accompanied this code.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

package jdk.internal.net.http;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.ListIterator;
import java.util.Objects;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.Flow;
import java.util.concurrent.atomic.AtomicBoolean;
import java.net.http.HttpResponse.BodySubscriber;
import jdk.internal.net.http.common.Demand;
import jdk.internal.net.http.common.SequentialScheduler;
import jdk.internal.net.http.ResponseSubscribers.TrustedSubscriber;

/**
 * A buffering BodySubscriber. When subscribed, accumulates ( buffers ) a given
 * amount ( in bytes ) of a publisher's data before pushing it to a downstream
 * subscriber.
 */
public class BufferingSubscriber<T> implements TrustedSubscriber<T>
{
    /** The downstream consumer of the data. */
    private final BodySubscriber<T> downstreamSubscriber;
    /** The amount of data to be accumulate before pushing downstream. */
    private final int bufferSize;

    /** The subscription, created lazily. */
    private volatile Flow.Subscription subscription;
    /** The downstream subscription, created lazily. */
    private volatile DownstreamSubscription downstreamSubscription;

    /** Must be held when accessing the internal buffers. */
    // using a monitor here is fine: no logging while holding it
    private final Object buffersLock = new Object();
    /** The internal buffers holding the buffered data. */
    private ArrayList<ByteBuffer> internalBuffers;
    /** The actual accumulated remaining bytes in internalBuffers. */
    private long accumulatedBytes;

    /** Holds the Throwable from upstream's onError. */
    private volatile Throwable throwable;

    /** State of the buffering subscriber:
     *  1) [UNSUBSCRIBED] when initially created
     *  2) [ACTIVE] when subscribed and can receive data
     *  3) [ERROR | CANCELLED | COMPLETE] (terminal state)
     */
    static final int UNSUBSCRIBED = 0x01;
    static final int ACTIVE       = 0x02;
    static final int ERROR        = 0x04;
    static final int CANCELLED    = 0x08;
    static final int COMPLETE     = 0x10;

    private volatile int state;

    public BufferingSubscriber(BodySubscriber<T> downstreamSubscriber,
                               int bufferSize) {
        this.downstreamSubscriber = Objects.requireNonNull(downstreamSubscriber);
        this.bufferSize = bufferSize;
        synchronized (buffersLock) {
            internalBuffers = new ArrayList<>();
        }
        state = UNSUBSCRIBED;
    }

    /** Returns the number of bytes remaining in the given buffers. */
    private static final long remaining(List<ByteBuffer> buffers) {
        return buffers.stream().mapToLong(ByteBuffer::remaining).sum();
    }

    @Override
    public boolean needsExecutor() {
        return TrustedSubscriber.needsExecutor(downstreamSubscriber);
    }

    /**
     * Tells whether, or not, there is at least a sufficient number of bytes
     * accumulated in the internal buffers. If the subscriber is COMPLETE, and
     * has some buffered data, then there is always enough ( to pass downstream ).
     */
    private final boolean hasEnoughAccumulatedBytes() {
        assert Thread.holdsLock(buffersLock);
        return accumulatedBytes >= bufferSize
                || (state == COMPLETE && accumulatedBytes > 0);
    }

    /**
     * Returns a new, unmodifiable, List<ByteBuffer> containing exactly the
     * amount of data as required before pushing downstream. The amount of data
     * may be less than required ( bufferSize ), in the case where the subscriber
     * is COMPLETE.
     */
    private List<ByteBuffer> fromInternalBuffers() {
        assert Thread.holdsLock(buffersLock);
        int leftToFill = bufferSize;
        int state = this.state;
        assert (state == ACTIVE || state == CANCELLED)
                ? accumulatedBytes >= leftToFill : true;
        List<ByteBuffer> dsts = new ArrayList<>();

        ListIterator<ByteBuffer> itr = internalBuffers.listIterator();
        while (itr.hasNext()) {
            ByteBuffer b = itr.next();
            if (b.remaining() <= leftToFill) {
                itr.remove();
                if (b.position() != 0)
                    b = b.slice();  // ensure position = 0 when propagated
                dsts.add(b);
                leftToFill -= b.remaining();
                accumulatedBytes -= b.remaining();
                if (leftToFill == 0)
                    break;
            } else {
                int prevLimit = b.limit();
                b.limit(b.position() + leftToFill);
                ByteBuffer slice = b.slice();
                dsts.add(slice);
                b.limit(prevLimit);
                b.position(b.position() + leftToFill);
                accumulatedBytes -= leftToFill;
                leftToFill = 0;
                break;
            }
        }
        assert (state == ACTIVE || state == CANCELLED)
                ? leftToFill == 0 : state == COMPLETE;
        assert (state == ACTIVE || state == CANCELLED)
                ? remaining(dsts) == bufferSize : state == COMPLETE;
        assert accumulatedBytes >= 0;
        assert dsts.stream().noneMatch(b -> b.position() != 0);
        return Collections.unmodifiableList(dsts);
    }

    /** Subscription that is passed to the downstream subscriber. */
    private class DownstreamSubscription implements Flow.Subscription {
        private final AtomicBoolean cancelled = new AtomicBoolean(); // false
        private final Demand demand = new Demand();
        private volatile boolean illegalArg;

        @Override
        public void request(long n) {
            if (cancelled.get() || illegalArg) {
                return;
            }
            if (n <= 0L) {
                // pass the "bad" value upstream so the Publisher can deal with
                // it appropriately, i.e. invoke onError
                illegalArg = true;
                subscription.request(n);
                return;
            }

            demand.increase(n);

            pushDemanded();
        }

        private final SequentialScheduler pushDemandedScheduler =
                new SequentialScheduler(new PushDemandedTask());

        void pushDemanded() {
            if (cancelled.get())
                return;
            pushDemandedScheduler.runOrSchedule();
        }

        class PushDemandedTask extends SequentialScheduler.CompleteRestartableTask {
            @Override
            public void run() {
                try {
                    Throwable t = throwable;
                    if (t != null) {
                        pushDemandedScheduler.stop(); // stop the demand scheduler
                        downstreamSubscriber.onError(t);
                        return;
                    }

                    while (true) {
                        List<ByteBuffer> item;
                        synchronized (buffersLock) {
                            if (cancelled.get())
                                return;
                            if (!hasEnoughAccumulatedBytes())
                                break;
                            if (!demand.tryDecrement())
                                break;
                            item = fromInternalBuffers();
                        }
                        assert item != null;

                        downstreamSubscriber.onNext(item);
                    }
                    if (cancelled.get())
                        return;

                    // complete only if all data consumed
                    boolean complete;
                    synchronized (buffersLock) {
                        complete = state == COMPLETE && internalBuffers.isEmpty();
                    }
                    if (complete) {
                        assert internalBuffers.isEmpty();
                        pushDemandedScheduler.stop(); // stop the demand scheduler
                        downstreamSubscriber.onComplete();
                        return;
                    }
                } catch (Throwable t) {
                    cancel();  // cancel if there is any error
                    throw t;
                }

                boolean requestMore = false;
                synchronized (buffersLock) {
                    if (!hasEnoughAccumulatedBytes() && !demand.isFulfilled()) {
                        // request more upstream data
                        requestMore = true;
                    }
                }
                if (requestMore)
                    subscription.request(1);
            }
        }

        @Override
        public void cancel() {
            if (cancelled.compareAndExchange(false, true))
                return;  // already cancelled

            state = CANCELLED;  // set CANCELLED state of upstream subscriber
            subscription.cancel();  // cancel upstream subscription
            pushDemandedScheduler.stop(); // stop the demand scheduler
        }
    }

    @Override
    public void onSubscribe(Flow.Subscription subscription) {
        Objects.requireNonNull(subscription);
        if (this.subscription != null) {
            subscription.cancel();
            return;
        }

        int s = this.state;
        assert s == UNSUBSCRIBED;
        state = ACTIVE;
        this.subscription = subscription;
        downstreamSubscription = new DownstreamSubscription();
        downstreamSubscriber.onSubscribe(downstreamSubscription);
    }

    @Override
    public void onNext(List<ByteBuffer> item) {
        Objects.requireNonNull(item);

        int s = state;
        if (s == CANCELLED)
            return;

        if (s != ACTIVE)
            throw new InternalError("onNext on inactive subscriber");

        synchronized (buffersLock) {
            internalBuffers.addAll(item);
            accumulatedBytes += remaining(item);
        }

        downstreamSubscription.pushDemanded();
    }

    @Override
    public void onError(Throwable incomingThrowable) {
        Objects.requireNonNull(incomingThrowable);
        int s = state;
        assert s == ACTIVE : "Expected ACTIVE, got:" + s;
        state = ERROR;
        Throwable t = this.throwable;
        assert t == null : "Expected null, got:" + t;
        this.throwable = incomingThrowable;
        downstreamSubscription.pushDemanded();
    }

    @Override
    public void onComplete() {
        int s = state;
        assert s == ACTIVE : "Expected ACTIVE, got:" + s;
        state = COMPLETE;
        downstreamSubscription.pushDemanded();
    }

    @Override
    public CompletionStage<T> getBody() {
        return downstreamSubscriber.getBody();
    }
}