<commit_before>/*
 * Copyright (c) 2008 - 2012 10gen, Inc. <http://10gen.com>
 *
 * 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 org.mongodb.io;

import org.bson.io.ByteBufferInput;
import org.bson.io.InputBuffer;
import org.bson.types.Document;
import org.bson.util.BufferPool;
import org.mongodb.MongoCursorNotFoundException;
import org.mongodb.MongoQueryFailureException;
import org.mongodb.ServerAddress;
import org.mongodb.protocol.MongoGetMoreMessage;
import org.mongodb.protocol.MongoQueryMessage;
import org.mongodb.protocol.MongoReplyHeader;
import org.mongodb.protocol.MongoReplyMessage;
import org.mongodb.protocol.MongoRequestMessage;
import org.mongodb.result.ServerCursor;
import org.mongodb.serialization.Serializer;

import java.io.IOException;
import java.io.InterruptedIOException;
import java.net.SocketTimeoutException;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;

import static org.mongodb.protocol.MongoReplyHeader.REPLY_HEADER_LENGTH;

// TODO: migrate all the DBPort configuration
// TODO: authentication

public class MongoChannel {
    private final ServerAddress address;
    private final BufferPool<ByteBuffer> pool;
    private final Serializer<Document> errorSerializer;
    private volatile SocketChannel socketChannel;

    public MongoChannel(final ServerAddress address, final BufferPool<ByteBuffer> pool,
                        final Serializer<Document> errorSerializer) {
        this.address = address;
        this.pool = pool;
        this.errorSerializer = errorSerializer;
        open();
    }

    public ServerAddress getAddress() {
        return address;
    }

    public void sendOneWayMessage(final MongoRequestMessage message) {

        try {
            if (socketChannel == null) {
                open();
            }

            message.pipe(socketChannel);
        } catch (IOException e) {
            throw new MongoSocketWriteException("Exception sending message", address, e);
        } finally {
            message.close();
        }
    }

    public <T> MongoReplyMessage<T> sendQueryMessage(final MongoQueryMessage message, final Serializer<T> serializer) {
        long start = System.nanoTime();
        sendOneWayMessage(message);
        return receiveMessage(message, serializer, start);
    }

    public <T> MongoReplyMessage<T> sendGetMoreMessage(final MongoGetMoreMessage message,
                                                       final Serializer<T> serializer) {
        long start = System.nanoTime();
        sendOneWayMessage(message);
        return receiveMessage(message, serializer, start);
    }

    private <T> MongoReplyMessage<T> receiveMessage(final MongoRequestMessage message, final Serializer<T> serializer, final long start) {
        ByteBuffer headerByteBuffer = null;
        ByteBuffer bodyByteBuffer = null;
        try {
            headerByteBuffer = pool.get(REPLY_HEADER_LENGTH);
            fillAndFlipBuffer(headerByteBuffer);
            final InputBuffer headerInputBuffer = new ByteBufferInput(headerByteBuffer);

            final MongoReplyHeader replyHeader = new MongoReplyHeader(headerInputBuffer);

            InputBuffer bodyInputBuffer = null;

            if (replyHeader.getNumberReturned() > 0) {
                bodyByteBuffer = pool.get(replyHeader.getMessageLength() - REPLY_HEADER_LENGTH);
                fillAndFlipBuffer(bodyByteBuffer);

                bodyInputBuffer = new ByteBufferInput(bodyByteBuffer);
            }

            if (replyHeader.isCursorNotFound()) {
                throw new MongoCursorNotFoundException(new ServerCursor(((MongoGetMoreMessage) message).getCursorId(),
                                                                       address));
            }
            else if (replyHeader.isQueryFailure()) {
                final Document errorDocument = new MongoReplyMessage<Document>(replyHeader, bodyInputBuffer,
                                                                              errorSerializer, System.nanoTime() - start).getDocuments().get(0);
                throw new MongoQueryFailureException(address, errorDocument);
            }
            return new MongoReplyMessage<T>(replyHeader, bodyInputBuffer, serializer, System.nanoTime() - start);
        } catch (SocketTimeoutException e) {
            throw new MongoSocketReadTimeoutException("Exception receiving message", address, e);
        } catch (InterruptedIOException e) {
            throw new MongoSocketInterruptedReadException("Exception receiving message", address, e);
        } catch (IOException e) {
            throw new MongoSocketReadException("Exception receiving message", address, e);
        } finally {
            if (headerByteBuffer != null) {
                pool.done(headerByteBuffer);
            }
            if (bodyByteBuffer != null) {
                pool.done(bodyByteBuffer);
            }
        }
    }

    private void fillAndFlipBuffer(final ByteBuffer buffer) throws IOException {
        int totalBytesRead = 0;
        while (totalBytesRead < buffer.limit()) {
            final int bytesRead = socketChannel.read(buffer);
            if (bytesRead == -1) {
                throw new MongoSocketReadException("Prematurely reached end of stream", address);
            }
            totalBytesRead += bytesRead;
        }
        buffer.flip();
    }

    private void open() {
        try {
            socketChannel = SocketChannel.open(address.getSocketAddress());<commit_msg>Changed location of call to open socket so that it's not included it time measurements<commit_after>/*
 * Copyright (c) 2008 - 2012 10gen, Inc. <http://10gen.com>
 *
 * 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 org.mongodb.io;

import org.bson.io.ByteBufferInput;
import org.bson.io.InputBuffer;
import org.bson.types.Document;
import org.bson.util.BufferPool;
import org.mongodb.MongoCursorNotFoundException;
import org.mongodb.MongoQueryFailureException;
import org.mongodb.ServerAddress;
import org.mongodb.protocol.MongoGetMoreMessage;
import org.mongodb.protocol.MongoQueryMessage;
import org.mongodb.protocol.MongoReplyHeader;
import org.mongodb.protocol.MongoReplyMessage;
import org.mongodb.protocol.MongoRequestMessage;
import org.mongodb.result.ServerCursor;
import org.mongodb.serialization.Serializer;

import java.io.IOException;
import java.io.InterruptedIOException;
import java.net.SocketTimeoutException;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;

import static org.mongodb.protocol.MongoReplyHeader.REPLY_HEADER_LENGTH;

// TODO: migrate all the DBPort configuration
// TODO: authentication

public class MongoChannel {
    private final ServerAddress address;
    private final BufferPool<ByteBuffer> pool;
    private final Serializer<Document> errorSerializer;
    private volatile SocketChannel socketChannel;

    public MongoChannel(final ServerAddress address, final BufferPool<ByteBuffer> pool,
                        final Serializer<Document> errorSerializer) {
        this.address = address;
        this.pool = pool;
        this.errorSerializer = errorSerializer;
    }

    public ServerAddress getAddress() {
        return address;
    }

    public void sendOneWayMessage(final MongoRequestMessage message) {
        try {
            ensureOpen();

            message.pipe(socketChannel);
        } catch (IOException e) {
            throw new MongoSocketWriteException("Exception sending message", address, e);
        } finally {
            message.close();
        }
    }

    public <T> MongoReplyMessage<T> sendQueryMessage(final MongoQueryMessage message, final Serializer<T> serializer) {
        ensureOpen();
        long start = System.nanoTime();
        sendOneWayMessage(message);
        return receiveMessage(message, serializer, start);
    }

    public <T> MongoReplyMessage<T> sendGetMoreMessage(final MongoGetMoreMessage message,
                                                       final Serializer<T> serializer) {
        ensureOpen();
        long start = System.nanoTime();
        sendOneWayMessage(message);
        return receiveMessage(message, serializer, start);
    }

    private <T> MongoReplyMessage<T> receiveMessage(final MongoRequestMessage message, final Serializer<T> serializer, final long start) {
        ByteBuffer headerByteBuffer = null;
        ByteBuffer bodyByteBuffer = null;
        try {
            headerByteBuffer = pool.get(REPLY_HEADER_LENGTH);
            fillAndFlipBuffer(headerByteBuffer);
            final InputBuffer headerInputBuffer = new ByteBufferInput(headerByteBuffer);

            final MongoReplyHeader replyHeader = new MongoReplyHeader(headerInputBuffer);

            InputBuffer bodyInputBuffer = null;

            if (replyHeader.getNumberReturned() > 0) {
                bodyByteBuffer = pool.get(replyHeader.getMessageLength() - REPLY_HEADER_LENGTH);
                fillAndFlipBuffer(bodyByteBuffer);

                bodyInputBuffer = new ByteBufferInput(bodyByteBuffer);
            }

            if (replyHeader.isCursorNotFound()) {
                throw new MongoCursorNotFoundException(new ServerCursor(((MongoGetMoreMessage) message).getCursorId(),
                        address));
            } else if (replyHeader.isQueryFailure()) {
                final Document errorDocument = new MongoReplyMessage<Document>(replyHeader, bodyInputBuffer,
                        errorSerializer, System.nanoTime() - start).getDocuments().get(0);
                throw new MongoQueryFailureException(address, errorDocument);
            }
            return new MongoReplyMessage<T>(replyHeader, bodyInputBuffer, serializer, System.nanoTime() - start);
        } catch (SocketTimeoutException e) {
            throw new MongoSocketReadTimeoutException("Exception receiving message", address, e);
        } catch (InterruptedIOException e) {
            throw new MongoSocketInterruptedReadException("Exception receiving message", address, e);
        } catch (IOException e) {
            throw new MongoSocketReadException("Exception receiving message", address, e);
        } finally {
            if (headerByteBuffer != null) {
                pool.done(headerByteBuffer);
            }
            if (bodyByteBuffer != null) {
                pool.done(bodyByteBuffer);
            }
        }
    }

    private void fillAndFlipBuffer(final ByteBuffer buffer) throws IOException {
        int totalBytesRead = 0;
        while (totalBytesRead < buffer.limit()) {
            final int bytesRead = socketChannel.read(buffer);
            if (bytesRead == -1) {
                throw new MongoSocketReadException("Prematurely reached end of stream", address);
            }
            totalBytesRead += bytesRead;
        }
        buffer.flip();
    }

    private void ensureOpen() {
        try {
            if (socketChannel == null) {
                socketChannel = SocketChannel.open(address.getSocketAddress());
            }
{'fraction_non_alphanumeric': 0.08113725164186548, 'fraction_numerical': 0.00332529719843711, 'mean_word_length': 2.7244582043343653, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 37, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4433820', 'n_tokens_mistral': 3040, 'n_tokens_neox': 2895, 'n_words': 788}