com.mongodb.QueryResultIterator 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
There is a newer version: 3.1.0
/*
 * Copyright (c) 2008-2014 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;

import org.bson.BSONObject;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;

import static com.mongodb.DBApiLayer.DeadCursor;
import static java.util.Arrays.asList;

class QueryResultIterator implements Cursor {

    private final DBDecoder _decoder;
    private final ServerAddress _host;
    private final int _limit;

    private DBApiLayer _db;
    private DBCollection _collection;
    private long _cursorId;
    private Iterator _cur;
    private int _curSize;
    private int _batchSize;

    private boolean closed;

    private final List _sizes = new ArrayList();
    private int _numGetMores = 0;
    private int _numFetched = 0;

    // This allows us to easily enable/disable finalizer for cleaning up un-closed cursors
    private OptionalFinalizer _optionalFinalizer;
    private boolean batchSizeTrackingDisabled;

    // Constructor to use for normal queries
    QueryResultIterator(String namespace, Mongo mongo, Response res, int batchSize, int limit, DBDecoder decoder) {
        _db = (DBApiLayer) mongo.getDB(getDatatabaseNameFromNamespace(namespace));
        _collection =  _db.getCollection(getCollectionNameFromNamespace(namespace));
        _batchSize = batchSize;
        _limit = limit;
        _host = res._host;
        _decoder = decoder;
        initFromQueryResponse(res, mongo);
    }

    // Constructor to use for commands that return cursor documents
    QueryResultIterator(DBObject cursorDocument, Mongo mongo, int batchSize, DBDecoder decoder, final ServerAddress serverAddress) {
        _batchSize = batchSize;
        _host = serverAddress;
        _limit = 0;
        _decoder = decoder;
        initFromCursorDocument(cursorDocument, mongo);
    }

    static int chooseBatchSize(int batchSize, int limit, int fetched) {
        int bs = Math.abs(batchSize);
        int remaining = limit > 0 ? limit - fetched : 0;
        int res;
        if (bs == 0 && remaining > 0)
            res = remaining;
        else if (bs > 0 && remaining == 0)
            res = bs;
        else
            res = Math.min(bs, remaining);

        if (batchSize < 0) {
            // force close
            res = -res;
        }

        if (res == 1) {
            // optimization: use negative batch size to close cursor
            res = -1;
        }
        return res;
    }

    public DBObject next() {
        if (closed) {
            throw new IllegalStateException("Iterator has been closed");
        }

        if (!hasNext()) {
            throw new NoSuchElementException();
        }

        return _cur.next();
    }

    public boolean tryHasNext() {
        if (closed) {
            throw new IllegalStateException("Iterator has been closed");
        }

        if (_cur.hasNext()) {
            return true;
        }

        if (_cursorId != 0) {
            getMore();
        }
        return _curSize > 0;

    }

    public boolean hasNext() {
        if (closed) {
           throw new IllegalStateException("Iterator has been closed");
        }

        if (_cur.hasNext()) {
            return true;
        }

        while (_cursorId != 0) {
            getMore();
            if (_curSize > 0) {
                return true;
            }
        }

        return false;
    }

    private void getMore(){
        Response res = _db._connector.call(_db, _collection,
                                           OutMessage.getMore(_collection, _cursorId, getGetMoreBatchSize()),
                                           _host, _decoder);
        _numGetMores++;
        initFromQueryResponse(res, _db.getMongo());
    }

    private int getGetMoreBatchSize() {
        return chooseBatchSize(_batchSize, _limit, _numFetched);
    }

    public void remove(){
        throw new UnsupportedOperationException("can't remove a document via a query result iterator");
    }

    public void setBatchSize(int size){
        _batchSize = size;
    }

    public long getCursorId(){
        return _cursorId;
    }

    int numGetMores(){
        return _numGetMores;
    }

    List getSizes(){
        return Collections.unmodifiableList(_sizes);
    }

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

    private void initFromQueryResponse(final Response response, final Mongo mongo) {
        init(response._flags, response.cursor(), response.size(), response.iterator(), mongo);
    }

    @SuppressWarnings("unchecked")
    private void initFromCursorDocument(final DBObject cursorDocument, final Mongo mongo) {
        Map cursor = (Map) cursorDocument.get("cursor");
        if (cursor != null) {
            long cursorId = (Long) cursor.get("id");
            List firstBatch = (List) cursor.get("firstBatch");
            String namespace = (String) cursor.get("ns");
            _db = (DBApiLayer) mongo.getDB(getDatatabaseNameFromNamespace(namespace));
            _collection = _db.getCollection(getCollectionNameFromNamespace(namespace));
            init(0, cursorId, firstBatch.size(), firstBatch.iterator(), mongo);
        } else {
            List result = (List) cursorDocument.get("result");
            init(0, 0, result.size(), result.iterator(), mongo);
        }
    }

    private String getCollectionNameFromNamespace(final String namespace) {
        return namespace.substring(namespace.indexOf('.') + 1);
    }

    private String getDatatabaseNameFromNamespace(final String namespace) {
        return namespace.substring(0, namespace.indexOf('.'));
    }

    private void init(int flags, long cursorId, int size, Iterator iterator, Mongo mongo){
        _curSize = size;
        _cur = iterator;
        if (!batchSizeTrackingDisabled) {
            _sizes.add(size);
        }
        _numFetched += size;

        if (_optionalFinalizer == null) {
            _optionalFinalizer = createFinalizerIfNeeded(cursorId, mongo);
        }

        setCursorIdOnFinalizer(cursorId);
        throwOnQueryFailure(_cursorId, flags);
        _cursorId = cursorId;


        if (cursorId != 0 && _limit > 0 && _limit - _numFetched <= 0) {
            // fetched all docs within limit, close cursor server-side
            killCursor();
        }
    }

    private void setCursorIdOnFinalizer(final long cursorId) {
        if (_optionalFinalizer != null) {
            _optionalFinalizer.setCursorId(cursorId);
        }
    }
    private void throwOnQueryFailure(final long cursorId, int flags) {
        if ((flags & Bytes.RESULTFLAG_ERRSET) > 0) {
            BSONObject errorDocument = _cur.next();
            if (ServerError.getCode(errorDocument) == 50) {
                throw new MongoExecutionTimeoutException(ServerError.getCode(errorDocument),
                                                         ServerError.getMsg(errorDocument, null));
            } else {
                throw new MongoException(ServerError.getCode(errorDocument), ServerError.getMsg(errorDocument, null));
            }
        }
        else if ((flags & Bytes.RESULTFLAG_CURSORNOTFOUND) > 0) {
            throw new MongoException.CursorNotFound(cursorId, _host);
        }
    }


    void killCursor() {
        setCursorIdOnFinalizer(0);

        if (_cursorId == 0)
            return;

        try {
            _db.killCursors(_host, asList(_cursorId));
            _cursorId = 0;
        } catch (MongoException e) {
            _db.addDeadCursor(new DeadCursor(_cursorId, _host));
        }
    }

    public ServerAddress getServerAddress() {
        return _host;
    }

    void disableBatchSizeTracking() {
        batchSizeTrackingDisabled = true;
        _sizes.clear();
    }

    boolean hasFinalizer() {
        return _optionalFinalizer != null;
    }

    private OptionalFinalizer createFinalizerIfNeeded(final long cursorId, Mongo mongo) {
        return mongo.getMongoOptions().isCursorFinalizerEnabled() && cursorId != 0 ?
               new OptionalFinalizer(_db, _host) : null;
    }

    private static class OptionalFinalizer {
        private final DBApiLayer db;
        private final ServerAddress host;
        private volatile long cursorId;

        private OptionalFinalizer(final DBApiLayer db, final ServerAddress host) {
            this.db = db;
            this.host = host;
        }

        public void setCursorId(final long cursorId) {
            this.cursorId = cursorId;
        }

        @Override
        protected void finalize() throws Throwable {
            if (cursorId != 0) {
                db.addDeadCursor(new DeadCursor(cursorId, host));
            }
        }
    }

}