com.mongodb.operation.AsyncQueryBatchCursor Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of mongo-java-driver Show documentation
The MongoDB Java Driver uber-artifact, containing mongodb-driver, mongodb-driver-core, and bson
The newest version!
/*
 * Copyright (c) 2008-2015 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.operation;

import com.mongodb.MongoNamespace;
import com.mongodb.ServerCursor;
import com.mongodb.async.AsyncBatchCursor;
import com.mongodb.async.SingleResultCallback;
import com.mongodb.binding.AsyncConnectionSource;
import com.mongodb.connection.AsyncConnection;
import com.mongodb.connection.QueryResult;
import org.bson.codecs.Decoder;

import java.util.List;

import static com.mongodb.assertions.Assertions.isTrue;
import static com.mongodb.assertions.Assertions.notNull;
import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
import static com.mongodb.operation.CursorHelper.getNumberToReturn;
import static java.util.Collections.singletonList;

class AsyncQueryBatchCursor implements AsyncBatchCursor {

    private final MongoNamespace namespace;
    private final int limit;
    private final Decoder decoder;
    private volatile AsyncConnectionSource connectionSource;
    private volatile QueryResult firstBatch;
    private volatile int batchSize;
    private volatile ServerCursor cursor;
    private volatile int count;
    private volatile boolean closed;

    AsyncQueryBatchCursor(final QueryResult firstBatch, final int limit, final int batchSize,
                          final Decoder decoder) {
        this(firstBatch, limit, batchSize, decoder, null, null);
    }

    AsyncQueryBatchCursor(final QueryResult firstBatch, final int limit, final int batchSize,
                          final Decoder decoder, final AsyncConnectionSource connectionSource, final AsyncConnection connection) {
        this.namespace = firstBatch.getNamespace();
        this.firstBatch = firstBatch;
        this.limit = limit;
        this.batchSize = batchSize;
        this.decoder = decoder;
        this.cursor = firstBatch.getCursor();
        if (this.cursor != null) {
            notNull("connectionSource", connectionSource);
            notNull("connection", connection);
        }
        if (connectionSource != null) {
            this.connectionSource = connectionSource.retain();
        } else {
            this.connectionSource = null;
        }
        this.count += firstBatch.getResults().size();
        if (limitReached()) {
             killCursor(connection);
        }
    }

    @Override
    public void close() {
        if (!closed) {
            closed = true;
            killCursor(null);
        }
    }

    @Override
    public void next(final SingleResultCallback> callback) {
        isTrue("open", !closed);
        // may be empty for a tailable cursor
        if (firstBatch  != null && !firstBatch.getResults().isEmpty()) {
            List results = firstBatch.getResults();
            firstBatch = null;
            callback.onResult(results, null);
        } else {
            if (cursor == null) {
                close();
                callback.onResult(null, null);
            } else {
                getMore(callback);
            }
        }
    }

    @Override
    public void setBatchSize(final int batchSize) {
        isTrue("open", !closed);
        this.batchSize = batchSize;
    }

    @Override
    public int getBatchSize() {
        isTrue("open", !closed);
        return batchSize;
    }

    @Override
    public boolean isClosed() {
        return closed;
    }

    private boolean limitReached() {
        return limit != 0 && count >= limit;
    }

    private void getMore(final SingleResultCallback> callback) {
        connectionSource.getConnection(new SingleResultCallback() {
            @Override
            public void onResult(final AsyncConnection connection, final Throwable t) {
                if (t != null) {
                    callback.onResult(null, t);
                } else {
                    connection.getMoreAsync(namespace, cursor.getId(), getNumberToReturn(limit, batchSize, count),
                                            decoder, new QueryResultSingleResultCallback(connection, callback));
                }
            }
        });
    }

    private void killCursor(final AsyncConnection connection) {
        if (cursor != null) {
            final ServerCursor localCursor = cursor;
            final AsyncConnectionSource localConnectionSource = connectionSource;
            cursor = null;
            connectionSource = null;
            if (connection != null) {
                connection.retain();
                killCursorAsynchronouslyAndReleaseConnectionAndSource(connection, localCursor, localConnectionSource);
            } else {
                localConnectionSource.getConnection(new SingleResultCallback() {
                    @Override
                    public void onResult(final AsyncConnection connection, final Throwable connectionException) {
                        if (connectionException == null) {
                            killCursorAsynchronouslyAndReleaseConnectionAndSource(connection, localCursor, localConnectionSource);
                        }
                    }
                });
            }
        } else if (connectionSource != null) {
            connectionSource.release();
            connectionSource = null;
        }
    }

    private void killCursorAsynchronouslyAndReleaseConnectionAndSource(final AsyncConnection connection, final ServerCursor localCursor,
                                                                       final AsyncConnectionSource localConnectionSource) {
        connection.killCursorAsync(namespace, singletonList(localCursor.getId()), new SingleResultCallback() {
            @Override
            public void onResult(final Void result, final Throwable t) {
                connection.release();
                localConnectionSource.release();
            }
        });
    }

    private class QueryResultSingleResultCallback implements SingleResultCallback> {
        private final AsyncConnection connection;
        private final SingleResultCallback> callback;

        public QueryResultSingleResultCallback(final AsyncConnection connection, final SingleResultCallback> callback) {
            this.connection = connection;
            this.callback = errorHandlingCallback(callback);
        }

        @Override
        public void onResult(final QueryResult result, final Throwable t) {
            if (t != null) {
                connection.release();
                close();
                callback.onResult(null, t);
            } else if (result.getResults().isEmpty() && result.getCursor() != null) {
                connection.getMoreAsync(namespace, cursor.getId(), getNumberToReturn(limit, batchSize, count),
                                        decoder, this);
            } else {
                cursor = result.getCursor();
                count += result.getResults().size();
                if (limitReached()) {
                    killCursor(connection);
                }
                connection.release();
                if (result.getResults().isEmpty()) {
                    callback.onResult(null, null);
                } else {
                    callback.onResult(result.getResults(), null);
                }
            }
        }
    }
}