io.smallrye.mutiny.operators.multi.MultiDemandCapping Maven / Gradle / Ivy

Go to download
package io.smallrye.mutiny.operators.multi;

import static io.smallrye.mutiny.helpers.ParameterValidation.nonNull;

import java.util.concurrent.Flow;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.LongFunction;

import io.smallrye.mutiny.Multi;
import io.smallrye.mutiny.helpers.Subscriptions;
import io.smallrye.mutiny.operators.MultiOperator;
import io.smallrye.mutiny.subscription.MultiSubscriber;

public class MultiDemandCapping extends MultiOperator {

    private final LongFunction function;

    public MultiDemandCapping(Multi upstream, LongFunction function) {
        super(upstream);
        this.function = function;
    }

    @Override
    public void subscribe(MultiSubscriber subscriber) {
        upstream().subscribe(new MultiDemandCappingProcessor(subscriber));
    }

    private class MultiDemandCappingProcessor extends MultiOperatorProcessor {

        private final AtomicLong demand = new AtomicLong();

        MultiDemandCappingProcessor(MultiSubscriber downstream) {
            super(downstream);
        }

        @Override
        public void request(long numberOfItems) {
            Flow.Subscription subscription = getUpstreamSubscription();
            if (subscription == Subscriptions.CANCELLED) {
                return;
            }
            if (numberOfItems <= 0) {
                onFailure(Subscriptions.getInvalidRequestException());
                return;
            }
            try {
                Subscriptions.add(demand, numberOfItems);
                long currentDemand = demand.get();
                long actualDemand = nonNull(function.apply(currentDemand), "actualDemand");
                if (actualDemand <= 0) {
                    onFailure(new IllegalArgumentException("Invalid number of request, must be greater than 0"));
                    return;
                }
                if (actualDemand > currentDemand) {
                    onFailure(new IllegalStateException("The demand capping function computed a request of " + actualDemand
                            + " elements while the outstanding demand is of " + numberOfItems + " elements"));
                    return;
                }
                Subscriptions.produced(demand, actualDemand);
                subscription.request(actualDemand);
            } catch (Throwable failure) {
                onFailure(failure);
            }
        }
    }
}