modules/network/src/main/java/org/apache/ignite/network/internal/netty/ConnectionManager.java - ignite-3 - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF 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.apache.ignite.network.internal.netty;

 import java.net.SocketAddress;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BiConsumer;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 import io.netty.bootstrap.Bootstrap;
 import io.netty.channel.ChannelOption;
 import io.netty.channel.EventLoopGroup;
 import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.channel.socket.nio.NioSocketChannel;
 import org.apache.ignite.lang.IgniteInternalException;
 import org.apache.ignite.lang.IgniteLogger;
 import org.apache.ignite.network.NetworkMessage;
 import org.apache.ignite.network.internal.handshake.HandshakeManager;
 import org.apache.ignite.network.serialization.MessageSerializationRegistry;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.TestOnly;

 /**
  * Class that manages connections both incoming and outgoing.
  */
 public class ConnectionManager {
     /** Logger. */
     private static final IgniteLogger LOG = IgniteLogger.forClass(ConnectionManager.class);

     /** Latest version of the direct marshalling protocol. */
     public static final byte DIRECT_PROTOCOL_VERSION = 1;

     /** Client bootstrap. */
     private final Bootstrap clientBootstrap;

     /** Client socket channel handler event loop group. */
     private final EventLoopGroup clientWorkerGroup = new NioEventLoopGroup();

     /** Server. */
     private final NettyServer server;

     /** Channels map from consistentId to {@link NettySender}. */
     private final Map<String, NettySender> channels = new ConcurrentHashMap<>();

     /** Clients. */
     private final Map<SocketAddress, NettyClient> clients = new ConcurrentHashMap<>();

     /** Serialization registry. */
     private final MessageSerializationRegistry serializationRegistry;

     /** Message listeners. */
     private final List<BiConsumer<SocketAddress, NetworkMessage>> listeners = new CopyOnWriteArrayList<>();

     /** Node consistent id. */
     private final String consistentId;

     /** Client handshake manager factory. */
     private final Supplier<HandshakeManager> clientHandshakeManagerFactory;

     /**
      * Constructor.
      *
      * @param port Server port.
      * @param registry Serialization registry.
      * @param consistentId Consistent id of this node.
      * @param serverHandshakeManagerFactory Server handshake manager factory.
      * @param clientHandshakeManagerFactory Client handshake manager factory.
      */
     public ConnectionManager(
         int port,
         MessageSerializationRegistry registry,
         String consistentId,
         Supplier<HandshakeManager> serverHandshakeManagerFactory,
         Supplier<HandshakeManager> clientHandshakeManagerFactory
     ) {
         this.serializationRegistry = registry;
         this.consistentId = consistentId;
         this.clientHandshakeManagerFactory = clientHandshakeManagerFactory;
         this.server = new NettyServer(
             port,
             serverHandshakeManagerFactory.get(),
             this::onNewIncomingChannel,
             this::onMessage,
             serializationRegistry
         );
         this.clientBootstrap = createClientBootstrap(clientWorkerGroup, serializationRegistry);
     }

     /**
      * Starts the server.
      *
      * @throws IgniteInternalException If failed to start.
      */
     public void start() throws IgniteInternalException {
         try {
             server.start().join();
         }
         catch (CompletionException e) {
             Throwable cause = e.getCause();
             throw new IgniteInternalException("Failed to start server: " + cause.getMessage(), cause);
         }
     }

     /**
      * @return Server local address.
      */
     public SocketAddress getLocalAddress() {
         return server.address();
     }

     /**
      * Gets a {@link NettySender}, that sends data from this node to another node with the specified address.
      * @param consistentId Another node's consistent id.
      * @param address Another node's address.
      * @return Sender.
      */
     public CompletableFuture<NettySender> channel(@Nullable String consistentId, SocketAddress address) {
         if (consistentId != null) {
             // If consistent id is known, try looking up a channel by consistent id. There can be an outbound connection
             // or an inbound connection associated with that consistent id.
             NettySender channel = channels.compute(
                 consistentId,
                 (addr, sender) -> (sender == null || !sender.isOpen()) ? null : sender
             );

             if (channel != null)
                 return CompletableFuture.completedFuture(channel);
         }

         // Get an existing client or create a new one. NettyClient provides a CompletableFuture that resolves
         // when the client is ready for write operations, so previously started client, that didn't establish connection
         // or didn't perform the handhsake operaton, can be reused.
         NettyClient client = clients.compute(address, (addr, existingClient) ->
             existingClient != null && !existingClient.failedToConnect() && !existingClient.isDisconnected() ?
                 existingClient : connect(addr)
         );

         CompletableFuture<NettySender> sender = client.sender();

         assert sender != null;

         return sender;
     }

     /**
      * Callback that is called upon receiving a new message.
      *
      * @param from Source of the message.
      * @param message New message.
      */
     private void onMessage(SocketAddress from, NetworkMessage message) {
         listeners.forEach(consumer -> consumer.accept(from, message));
     }

     /**
      * Callback that is called upon new client connected to the server.
      *
      * @param channel Channel from client to this {@link #server}.
      */
     private void onNewIncomingChannel(NettySender channel) {
         channels.put(channel.consistentId(), channel);
     }

     /**
      * Create new client from this node to specified address.
      *
      * @param address Target address.
      * @return New netty client.
      */
     private NettyClient connect(
         SocketAddress address
     ) {
         var client = new NettyClient(
             address,
             serializationRegistry,
             clientHandshakeManagerFactory.get(),
             this::onMessage
         );

         client.start(clientBootstrap).whenComplete((sender, throwable) -> {
             if (throwable == null)
                 channels.put(sender.consistentId(), sender);
             else
                 clients.remove(address);
         });

         return client;
     }

     /**
      * Add incoming message listener.
      *
      * @param listener Message listener.
      */
     public void addListener(BiConsumer<SocketAddress, NetworkMessage> listener) {
         listeners.add(listener);
     }

     /**
      * Stops the server and all clients.
      */
     public void stop() {
          Stream<CompletableFuture<Void>> stream = Stream.concat(
             clients.values().stream().map(NettyClient::stop),
             Stream.of(server.stop())
         );

          CompletableFuture<Void> stopFut = CompletableFuture.allOf(stream.toArray(CompletableFuture<?>[]::new));

          try {
              stopFut.join();
              // TODO: IGNITE-14538 quietPeriod and timeout should be configurable.
              clientWorkerGroup.shutdownGracefully(0L, 15, TimeUnit.SECONDS).sync();
          }
          catch (Exception e) {
              LOG.warn("Failed to stop the ConnectionManager: " + e.getMessage());
          }
     }

     /**
      * @return Connection manager's {@link #server}.
      */
     @TestOnly
     public NettyServer server() {
         return server;
     }

     /**
      * @return This node's consistent id.
      */
     @TestOnly
     public String consistentId() {
         return consistentId;
     }

     /**
      * @return Collection of all the clients started by this connection manager.
      */
     @TestOnly
     public Collection<NettyClient> clients() {
         return Collections.unmodifiableCollection(clients.values());
     }


     /**
      * @return Map of the channels.
      */
     @TestOnly
     public Map<String, NettySender> channels() {
         return Collections.unmodifiableMap(channels);
     }

     /**
      * Creates a {@link Bootstrap} for clients, providing channel handlers and options.
      *
      * @param eventLoopGroup Event loop group for channel handling.
      * @param serializationRegistry Serialization registry.
      * @return Bootstrap for clients.
      */
     public static Bootstrap createClientBootstrap(
         EventLoopGroup eventLoopGroup,
         MessageSerializationRegistry serializationRegistry
     ) {
         Bootstrap clientBootstrap = new Bootstrap();

         clientBootstrap.group(eventLoopGroup)
             .channel(NioSocketChannel.class)
             // See NettyServer#start for netty configuration details.
             .option(ChannelOption.SO_KEEPALIVE, true)
             .option(ChannelOption.SO_LINGER, 0)
             .option(ChannelOption.TCP_NODELAY, true);

         return clientBootstrap;
     }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF 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.apache.ignite.network.internal.netty;

	import java.net.SocketAddress;
	import java.util.Collection;
	import java.util.Collections;
	import java.util.List;
	import java.util.Map;
	import java.util.concurrent.CompletableFuture;
	import java.util.concurrent.CompletionException;
	import java.util.concurrent.ConcurrentHashMap;
	import java.util.concurrent.CopyOnWriteArrayList;
	import java.util.concurrent.TimeUnit;
	import java.util.function.BiConsumer;
	import java.util.function.Supplier;
	import java.util.stream.Stream;
	import io.netty.bootstrap.Bootstrap;
	import io.netty.channel.ChannelOption;
	import io.netty.channel.EventLoopGroup;
	import io.netty.channel.nio.NioEventLoopGroup;
	import io.netty.channel.socket.nio.NioSocketChannel;
	import org.apache.ignite.lang.IgniteInternalException;
	import org.apache.ignite.lang.IgniteLogger;
	import org.apache.ignite.network.NetworkMessage;
	import org.apache.ignite.network.internal.handshake.HandshakeManager;
	import org.apache.ignite.network.serialization.MessageSerializationRegistry;
	import org.jetbrains.annotations.Nullable;
	import org.jetbrains.annotations.TestOnly;

	/**
	* Class that manages connections both incoming and outgoing.
	*/
	public class ConnectionManager {
	/** Logger. */
	private static final IgniteLogger LOG = IgniteLogger.forClass(ConnectionManager.class);

	/** Latest version of the direct marshalling protocol. */
	public static final byte DIRECT_PROTOCOL_VERSION = 1;

	/** Client bootstrap. */
	private final Bootstrap clientBootstrap;

	/** Client socket channel handler event loop group. */
	private final EventLoopGroup clientWorkerGroup = new NioEventLoopGroup();

	/** Server. */
	private final NettyServer server;

	/** Channels map from consistentId to {@link NettySender}. */
	private final Map<String, NettySender> channels = new ConcurrentHashMap<>();

	/** Clients. */
	private final Map<SocketAddress, NettyClient> clients = new ConcurrentHashMap<>();

	/** Serialization registry. */
	private final MessageSerializationRegistry serializationRegistry;

	/** Message listeners. */
	private final List<BiConsumer<SocketAddress, NetworkMessage>> listeners = new CopyOnWriteArrayList<>();

	/** Node consistent id. */
	private final String consistentId;

	/** Client handshake manager factory. */
	private final Supplier<HandshakeManager> clientHandshakeManagerFactory;

	/**
	* Constructor.
	*
	* @param port Server port.
	* @param registry Serialization registry.
	* @param consistentId Consistent id of this node.
	* @param serverHandshakeManagerFactory Server handshake manager factory.
	* @param clientHandshakeManagerFactory Client handshake manager factory.
	*/
	public ConnectionManager(
	int port,
	MessageSerializationRegistry registry,
	String consistentId,
	Supplier<HandshakeManager> serverHandshakeManagerFactory,
	Supplier<HandshakeManager> clientHandshakeManagerFactory
	) {
	this.serializationRegistry = registry;
	this.consistentId = consistentId;
	this.clientHandshakeManagerFactory = clientHandshakeManagerFactory;
	this.server = new NettyServer(
	port,
	serverHandshakeManagerFactory.get(),
	this::onNewIncomingChannel,
	this::onMessage,
	serializationRegistry
	);
	this.clientBootstrap = createClientBootstrap(clientWorkerGroup, serializationRegistry);
	}

	/**
	* Starts the server.
	*
	* @throws IgniteInternalException If failed to start.
	*/
	public void start() throws IgniteInternalException {
	try {
	server.start().join();
	}
	catch (CompletionException e) {
	Throwable cause = e.getCause();
	throw new IgniteInternalException("Failed to start server: " + cause.getMessage(), cause);
	}
	}

	/**
	* @return Server local address.
	*/
	public SocketAddress getLocalAddress() {
	return server.address();
	}

	/**
	* Gets a {@link NettySender}, that sends data from this node to another node with the specified address.
	* @param consistentId Another node's consistent id.
	* @param address Another node's address.
	* @return Sender.
	*/
	public CompletableFuture<NettySender> channel(@Nullable String consistentId, SocketAddress address) {
	if (consistentId != null) {
	// If consistent id is known, try looking up a channel by consistent id. There can be an outbound connection
	// or an inbound connection associated with that consistent id.
	NettySender channel = channels.compute(
	consistentId,
	(addr, sender) -> (sender == null \|\| !sender.isOpen()) ? null : sender
	);

	if (channel != null)
	return CompletableFuture.completedFuture(channel);
	}

	// Get an existing client or create a new one. NettyClient provides a CompletableFuture that resolves
	// when the client is ready for write operations, so previously started client, that didn't establish connection
	// or didn't perform the handhsake operaton, can be reused.
	NettyClient client = clients.compute(address, (addr, existingClient) ->
	existingClient != null && !existingClient.failedToConnect() && !existingClient.isDisconnected() ?
	existingClient : connect(addr)
	);

	CompletableFuture<NettySender> sender = client.sender();

	assert sender != null;

	return sender;
	}

	/**
	* Callback that is called upon receiving a new message.
	*
	* @param from Source of the message.
	* @param message New message.
	*/
	private void onMessage(SocketAddress from, NetworkMessage message) {
	listeners.forEach(consumer -> consumer.accept(from, message));
	}

	/**
	* Callback that is called upon new client connected to the server.
	*
	* @param channel Channel from client to this {@link #server}.
	*/
	private void onNewIncomingChannel(NettySender channel) {
	channels.put(channel.consistentId(), channel);
	}

	/**
	* Create new client from this node to specified address.
	*
	* @param address Target address.
	* @return New netty client.
	*/
	private NettyClient connect(
	SocketAddress address
	) {
	var client = new NettyClient(
	address,
	serializationRegistry,
	clientHandshakeManagerFactory.get(),
	this::onMessage
	);

	client.start(clientBootstrap).whenComplete((sender, throwable) -> {
	if (throwable == null)
	channels.put(sender.consistentId(), sender);
	else
	clients.remove(address);
	});

	return client;
	}

	/**
	* Add incoming message listener.
	*
	* @param listener Message listener.
	*/
	public void addListener(BiConsumer<SocketAddress, NetworkMessage> listener) {
	listeners.add(listener);
	}

	/**
	* Stops the server and all clients.
	*/
	public void stop() {
	Stream<CompletableFuture<Void>> stream = Stream.concat(
	clients.values().stream().map(NettyClient::stop),
	Stream.of(server.stop())
	);

	CompletableFuture<Void> stopFut = CompletableFuture.allOf(stream.toArray(CompletableFuture<?>[]::new));

	try {
	stopFut.join();
	// TODO: IGNITE-14538 quietPeriod and timeout should be configurable.
	clientWorkerGroup.shutdownGracefully(0L, 15, TimeUnit.SECONDS).sync();
	}
	catch (Exception e) {
	LOG.warn("Failed to stop the ConnectionManager: " + e.getMessage());
	}
	}

	/**
	* @return Connection manager's {@link #server}.
	*/
	@TestOnly
	public NettyServer server() {
	return server;
	}

	/**
	* @return This node's consistent id.
	*/
	@TestOnly
	public String consistentId() {
	return consistentId;
	}

	/**
	* @return Collection of all the clients started by this connection manager.
	*/
	@TestOnly
	public Collection<NettyClient> clients() {
	return Collections.unmodifiableCollection(clients.values());
	}


	/**
	* @return Map of the channels.
	*/
	@TestOnly
	public Map<String, NettySender> channels() {
	return Collections.unmodifiableMap(channels);
	}

	/**
	* Creates a {@link Bootstrap} for clients, providing channel handlers and options.
	*
	* @param eventLoopGroup Event loop group for channel handling.
	* @param serializationRegistry Serialization registry.
	* @return Bootstrap for clients.
	*/
	public static Bootstrap createClientBootstrap(
	EventLoopGroup eventLoopGroup,
	MessageSerializationRegistry serializationRegistry
	) {
	Bootstrap clientBootstrap = new Bootstrap();

	clientBootstrap.group(eventLoopGroup)
	.channel(NioSocketChannel.class)
	// See NettyServer#start for netty configuration details.
	.option(ChannelOption.SO_KEEPALIVE, true)
	.option(ChannelOption.SO_LINGER, 0)
	.option(ChannelOption.TCP_NODELAY, true);

	return clientBootstrap;
	}
	}