<commit_before>/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch licenses this file to you 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.elasticsearch.transport;

import org.elasticsearch.Version;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.cli.SuppressForbidden;
import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.io.stream.BytesStreamOutput;
import org.elasticsearch.common.io.stream.InputStreamStreamInput;
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
import org.elasticsearch.common.io.stream.ReleasableBytesStreamOutput;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.network.NetworkService;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.util.BigArrays;
import org.elasticsearch.common.util.CancellableThreads;
import org.elasticsearch.common.util.concurrent.AbstractRunnable;
import org.elasticsearch.common.util.concurrent.EsExecutors;
import io.crate.common.io.IOUtils;
import io.crate.concurrent.CompletableContext;

import org.elasticsearch.indices.breaker.CircuitBreakerService;
import org.elasticsearch.threadpool.ThreadPool;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketException;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * This is a socket based blocking TcpTransport implementation that is used for tests
 * that need real networking. This implementation is a test only implementation that implements
 * the networking layer in the worst possible way since it blocks and uses a thread per request model.
 */
public class MockTcpTransport extends TcpTransport {

    /**
     * A pre-built light connection profile that shares a single connection across all
     * types.
     */
    public static final ConnectionProfile LIGHT_PROFILE;

    private final Set<MockChannel> openChannels = new HashSet<>();

    static {
        ConnectionProfile.Builder builder = new ConnectionProfile.Builder();
        builder.addConnections(1,
                               TransportRequestOptions.Type.BULK,
                               TransportRequestOptions.Type.PING,
                               TransportRequestOptions.Type.RECOVERY,
                               TransportRequestOptions.Type.REG,
                               TransportRequestOptions.Type.STATE);
        LIGHT_PROFILE = builder.build();
    }

    private final ExecutorService executor;

    public MockTcpTransport(Settings settings, ThreadPool threadPool, BigArrays bigArrays,
                            CircuitBreakerService circuitBreakerService, NamedWriteableRegistry namedWriteableRegistry,
                            NetworkService networkService) {
        this(settings, threadPool, bigArrays, circuitBreakerService, namedWriteableRegistry, networkService,
             Version.CURRENT);
    }

    public MockTcpTransport(Settings settings, ThreadPool threadPool, BigArrays bigArrays,
                            CircuitBreakerService circuitBreakerService, NamedWriteableRegistry namedWriteableRegistry,
                            NetworkService networkService, Version mockVersion) {
        super("mock-tcp-transport",
              settings,
              mockVersion,
              threadPool,
              bigArrays,
              circuitBreakerService,
              namedWriteableRegistry,
              networkService);
        // we have our own crazy cached threadpool this one is not bounded at all...
        // using the ES thread factory here is crucial for tests otherwise disruption tests won't block that thread
        executor = Executors.newCachedThreadPool(EsExecutors.daemonThreadFactory(settings,
                                                                                 Transports.TEST_MOCK_TRANSPORT_THREAD_PREFIX));
    }

    @Override
    protected MockChannel bind(final String name, InetSocketAddress address) throws IOException {
        ServerSocket socket = new ServerSocket();
        socket.setReuseAddress(TransportSettings.TCP_REUSE_ADDRESS.get(settings));
        ByteSizeValue tcpReceiveBufferSize = TransportSettings.TCP_RECEIVE_BUFFER_SIZE.get(settings);
        if (tcpReceiveBufferSize.getBytes() > 0) {
            socket.setReceiveBufferSize(tcpReceiveBufferSize.bytesAsInt());
        }
        socket.bind(address);
        MockChannel serverMockChannel = new MockChannel(socket, name);
        CountDownLatch started = new CountDownLatch(1);
        executor.execute(new AbstractRunnable() {
            @Override
            public void onFailure(Exception e) {
                onException(serverMockChannel, e);
            }

            @Override
            protected void doRun() throws Exception {
                started.countDown();
                serverMockChannel.accept(executor);
            }
        });
        try {
            started.await();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
        return serverMockChannel;
    }

    private void readMessage(MockChannel mockChannel, StreamInput input) throws IOException {
        Socket socket = mockChannel.activeChannel;
        byte[] minimalHeader = new byte[TcpHeader.MARKER_BYTES_SIZE];
        int firstByte = input.read();
        if (firstByte == -1) {
            throw new IOException("Connection reset by peer");
        }
        minimalHeader[0] = (byte) firstByte;
        minimalHeader[1] = (byte) input.read();
        int msgSize = input.readInt();
        if (msgSize == -1) {
            socket.getOutputStream().flush();
        } else {
            try (BytesStreamOutput output = new ReleasableBytesStreamOutput(msgSize, bigArrays)) {
                final byte[] buffer = new byte[msgSize];
                input.readFully(buffer);
                output.write(minimalHeader);
                output.writeInt(msgSize);
                output.write(buffer);
                consumeNetworkReads(mockChannel, output.bytes());
            }<commit_msg>bp: Improve MockTcpTransport memory usage<commit_after>/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch licenses this file to you 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.elasticsearch.transport;

import org.elasticsearch.Version;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.cli.SuppressForbidden;
import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.common.bytes.BytesArray;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.io.stream.BytesStreamOutput;
import org.elasticsearch.common.io.stream.InputStreamStreamInput;
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
import org.elasticsearch.common.io.stream.ReleasableBytesStreamOutput;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.network.NetworkService;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.util.BigArrays;
import org.elasticsearch.common.util.CancellableThreads;
import org.elasticsearch.common.util.concurrent.AbstractRunnable;
import org.elasticsearch.common.util.concurrent.EsExecutors;
import io.crate.common.io.IOUtils;
import io.crate.concurrent.CompletableContext;

import org.elasticsearch.indices.breaker.CircuitBreakerService;
import org.elasticsearch.threadpool.ThreadPool;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.Closeable;
import java.io.EOFException;
import java.io.IOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketException;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * This is a socket based blocking TcpTransport implementation that is used for tests
 * that need real networking. This implementation is a test only implementation that implements
 * the networking layer in the worst possible way since it blocks and uses a thread per request model.
 */
public class MockTcpTransport extends TcpTransport {

    /**
     * A pre-built light connection profile that shares a single connection across all
     * types.
     */
    public static final ConnectionProfile LIGHT_PROFILE;

    private final Set<MockChannel> openChannels = new HashSet<>();

    static {
        ConnectionProfile.Builder builder = new ConnectionProfile.Builder();
        builder.addConnections(1,
                               TransportRequestOptions.Type.BULK,
                               TransportRequestOptions.Type.PING,
                               TransportRequestOptions.Type.RECOVERY,
                               TransportRequestOptions.Type.REG,
                               TransportRequestOptions.Type.STATE);
        LIGHT_PROFILE = builder.build();
    }

    private final ExecutorService executor;

    public MockTcpTransport(Settings settings, ThreadPool threadPool, BigArrays bigArrays,
                            CircuitBreakerService circuitBreakerService, NamedWriteableRegistry namedWriteableRegistry,
                            NetworkService networkService) {
        this(settings, threadPool, bigArrays, circuitBreakerService, namedWriteableRegistry, networkService,
             Version.CURRENT);
    }

    public MockTcpTransport(Settings settings, ThreadPool threadPool, BigArrays bigArrays,
                            CircuitBreakerService circuitBreakerService, NamedWriteableRegistry namedWriteableRegistry,
                            NetworkService networkService, Version mockVersion) {
        super("mock-tcp-transport",
              settings,
              mockVersion,
              threadPool,
              bigArrays,
              circuitBreakerService,
              namedWriteableRegistry,
              networkService);
        // we have our own crazy cached threadpool this one is not bounded at all...
        // using the ES thread factory here is crucial for tests otherwise disruption tests won't block that thread
        executor = Executors.newCachedThreadPool(EsExecutors.daemonThreadFactory(settings,
                                                                                 Transports.TEST_MOCK_TRANSPORT_THREAD_PREFIX));
    }

    @Override
    protected MockChannel bind(final String name, InetSocketAddress address) throws IOException {
        ServerSocket socket = new ServerSocket();
        socket.setReuseAddress(TransportSettings.TCP_REUSE_ADDRESS.get(settings));
        ByteSizeValue tcpReceiveBufferSize = TransportSettings.TCP_RECEIVE_BUFFER_SIZE.get(settings);
        if (tcpReceiveBufferSize.getBytes() > 0) {
            socket.setReceiveBufferSize(tcpReceiveBufferSize.bytesAsInt());
        }
        socket.bind(address);
        MockChannel serverMockChannel = new MockChannel(socket, name);
        CountDownLatch started = new CountDownLatch(1);
        executor.execute(new AbstractRunnable() {
            @Override
            public void onFailure(Exception e) {
                onException(serverMockChannel, e);
            }

            @Override
            protected void doRun() throws Exception {
                started.countDown();
                serverMockChannel.accept(executor);
            }
        });
        try {
            started.await();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
        return serverMockChannel;
    }

    private void readMessage(MockChannel mockChannel, StreamInput input) throws IOException {
        Socket socket = mockChannel.activeChannel;
        byte[] minimalHeader = new byte[TcpHeader.MARKER_BYTES_SIZE + TcpHeader.MESSAGE_LENGTH_SIZE];
        try {
            input.readFully(minimalHeader);
        } catch (EOFException eof) {
            throw new IOException("Connection reset by peer");
        }

        // Read message length will throw stream corrupted exception if the marker bytes incorrect
        int msgSize = TcpTransport.readMessageLength(new BytesArray(minimalHeader));
        if (msgSize == -1) {
            socket.getOutputStream().flush();
        } else {
            final byte[] buffer = new byte[msgSize];
            input.readFully(buffer);
            int expectedSize = TcpHeader.MARKER_BYTES_SIZE + TcpHeader.MESSAGE_LENGTH_SIZE + msgSize;
            try (BytesStreamOutput output = new ReleasableBytesStreamOutput(expectedSize, bigArrays)) {
                output.write(minimalHeader);
                output.write(buffer);
                consumeNetworkReads(mockChannel, output.bytes());
            }
{'fraction_non_alphanumeric': 0.07323048400979938, 'fraction_numerical': 0.001258028206316626, 'mean_word_length': 2.9716013673415724, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '140216', 'n_tokens_mistral': 3609, 'n_tokens_neox': 3618, 'n_words': 968}