com.mongodb.reactivestreams.client.internal.BatchCursorFlux Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of mongodb-driver-reactivestreams Show documentation
A Reactive Streams implementation of the MongoDB Java driver
There is a newer version: 5.3.0-beta0
Show newest version
/*
 * Copyright 2008-present MongoDB, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.mongodb.reactivestreams.client.internal;

import org.reactivestreams.Publisher;
import org.reactivestreams.Subscriber;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;
import reactor.core.publisher.Mono;

import java.util.Objects;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

class BatchCursorFlux implements Publisher {

    private final BatchCursorPublisher batchCursorPublisher;
    private final AtomicBoolean inProgress = new AtomicBoolean(false);
    private final AtomicLong demandDelta = new AtomicLong(0);
    private volatile BatchCursor batchCursor;
    private FluxSink sink;

    BatchCursorFlux(final BatchCursorPublisher batchCursorPublisher) {
        this.batchCursorPublisher = batchCursorPublisher;
    }

    @Override
    public void subscribe(final Subscriber subscriber) {
        Flux.create(sink -> {
            this.sink = sink;
            sink.onRequest(demand -> {
                if (calculateDemand(demand) > 0 && inProgress.compareAndSet(false, true)) {
                    if (batchCursor == null) {
                        int batchSize = calculateBatchSize(sink.requestedFromDownstream());
                        batchCursorPublisher.batchCursor(batchSize)
                                .contextWrite(sink.contextView())
                                .subscribe(bc -> {
                            batchCursor = bc;
                            inProgress.set(false);

                            // Handle any cancelled subscriptions that happen during the time it takes to get the batchCursor
                            if (sink.isCancelled()) {
                                closeCursor();
                            } else {
                                recurseCursor();
                            }
                        }, sink::error);
                    } else {
                        inProgress.set(false);
                        recurseCursor();
                    }
                }
            });
            sink.onCancel(this::closeCursor);
            sink.onDispose(this::closeCursor);
        }, FluxSink.OverflowStrategy.BUFFER)
        .subscribe(subscriber);
    }

    private void closeCursor() {
        if (batchCursor != null) {
            batchCursor.close();
        }
    }

    private void recurseCursor(){
        if (!sink.isCancelled() && sink.requestedFromDownstream() > 0 && inProgress.compareAndSet(false, true)) {
            if (batchCursor.isClosed()) {
                sink.complete();
            } else {
                batchCursor.setBatchSize(calculateBatchSize(sink.requestedFromDownstream()));
                Mono.from(batchCursor.next(() -> sink.isCancelled()))
                        .contextWrite(sink.contextView())
                        .doOnCancel(this::closeCursor)
                        .subscribe(results -> {
                                    if (!results.isEmpty()) {
                                        results
                                                .stream()
                                                .filter(Objects::nonNull)
                                                .forEach(sink::next);
                                        calculateDemand(-results.size());
                                    }
                                    if (batchCursor.isClosed()) {
                                        sink.complete();
                                    } else {
                                        inProgress.set(false);
                                        recurseCursor();
                                    }
                                },
                                e -> {
                                    try {
                                        closeCursor();
                                    } finally {
                                        sink.error(e);
                                    }
                                });
                }
        }
    }

    long calculateDemand(final long demand) {
        return demandDelta.accumulateAndGet(demand, (originalValue, update) -> {
            long newValue = originalValue + update;
            return update > 0 && newValue < originalValue ? Long.MAX_VALUE : newValue;
        });
    }

    int calculateBatchSize(final long demand) {
        Integer setBatchSize = batchCursorPublisher.getBatchSize();
        if (setBatchSize != null) {
            return setBatchSize;
        } else if (demand > Integer.MAX_VALUE) {
            return Integer.MAX_VALUE;
        }
        return Math.max(2, (int) demand);
    }

}