IGNITE-27414 : Use MessageSerializer for TcpDiscoveryClientReconnectMessage v2

modules/core/src/main/java/org/apache/ignite/internal/managers/discovery/DiscoveryMessageFactory.java

@@ -17,10 +17,17 @@
 
 package org.apache.ignite.internal.managers.discovery;
 
+import java.util.function.Supplier;
+import org.apache.ignite.IgniteException;
 import org.apache.ignite.internal.processors.cache.CacheStatisticsModeChangeMessage;
 import org.apache.ignite.internal.processors.cache.CacheStatisticsModeChangeMessageSerializer;
+import org.apache.ignite.marshaller.Marshaller;
+import org.apache.ignite.plugin.extensions.communication.Message;
 import org.apache.ignite.plugin.extensions.communication.MessageFactory;
 import org.apache.ignite.plugin.extensions.communication.MessageFactoryProvider;
+import org.apache.ignite.plugin.extensions.communication.MessageReader;
+import org.apache.ignite.plugin.extensions.communication.MessageSerializer;
+import org.apache.ignite.plugin.extensions.communication.MessageWriter;
 import org.apache.ignite.spi.discovery.tcp.internal.DiscoveryDataPacket;
 import org.apache.ignite.spi.discovery.tcp.internal.DiscoveryDataPacketSerializer;
 import org.apache.ignite.spi.discovery.tcp.messages.InetAddressMessage;
@@ -45,6 +52,10 @@
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryClientPingRequestSerializer;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryClientPingResponse;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryClientPingResponseSerializer;
+import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryClientReconnectMessage;
+import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryClientReconnectMessageSerializer;
+import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryCollectionMessage;
+import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryCollectionMessageSerializer;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryConnectionCheckMessage;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryConnectionCheckMessageSerializer;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryCustomEventMessage;
@@ -61,6 +72,7 @@
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryJoinRequestMessageSerializer;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryLoopbackProblemMessage;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryLoopbackProblemMessageSerializer;
+import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryMarshallableMessage;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryMetricsUpdateMessage;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryMetricsUpdateMessageSerializer;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryNodeAddFinishedMessage;
@@ -83,11 +95,35 @@
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryServerOnlyCustomEventMessageSerializer;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryStatusCheckMessage;
 import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryStatusCheckMessageSerializer;
+import org.jetbrains.annotations.Nullable;
 
-/** Message factory for discovery messages. */
+/**
+ * Message factory for discovery messages. Allows to create an enhanced {@link MessageFactory} allowing to create
+ * automated pre- and post- marshalling message serializer for {@link TcpDiscoveryMarshallableMessage}.
+ */
 public class DiscoveryMessageFactory implements MessageFactoryProvider {
+    /** Custom data marshaller. */
+    private final @Nullable Marshaller cstDataMarshall;
+
+    /** Class loader for the custom data marshalling. */
+    private final @Nullable ClassLoader cstDataMarshallClsLdr;
+
+    /**
+     * @param cstDataMarshall Custom data marshaller.
+     * @param cstDataMarshallClsLdr Class loader for the custom data marshalling.
+     */
+    public DiscoveryMessageFactory(@Nullable Marshaller cstDataMarshall, @Nullable ClassLoader cstDataMarshallClsLdr) {
+        assert cstDataMarshall == null && cstDataMarshallClsLdr == null || cstDataMarshall != null && cstDataMarshallClsLdr != null;
+
+        this.cstDataMarshall = cstDataMarshall;
+        this.cstDataMarshallClsLdr = cstDataMarshallClsLdr;
+    }
+
     /** {@inheritDoc} */
     @Override public void registerAll(MessageFactory factory) {
+        factory = enhanceMessageFactory(factory);
+
+        factory.register((short)-108, TcpDiscoveryCollectionMessage::new, new TcpDiscoveryCollectionMessageSerializer());
         factory.register((short)-107, NodeSpecificData::new, new NodeSpecificDataSerializer());
         factory.register((short)-106, DiscoveryDataPacket::new, new DiscoveryDataPacketSerializer());
         factory.register((short)-105, TcpDiscoveryNodeFullMetricsMessage::new,
@@ -123,8 +159,75 @@ public class DiscoveryMessageFactory implements MessageFactoryProvider {
         factory.register((short)21, TcpDiscoveryCustomEventMessage::new, new TcpDiscoveryCustomEventMessageSerializer());
         factory.register((short)22, TcpDiscoveryServerOnlyCustomEventMessage::new,
             new TcpDiscoveryServerOnlyCustomEventMessageSerializer());
+        factory.register((short)23, TcpDiscoveryClientReconnectMessage::new, new TcpDiscoveryClientReconnectMessageSerializer());
 
         // DiscoveryCustomMessage
         factory.register((short)500, CacheStatisticsModeChangeMessage::new, new CacheStatisticsModeChangeMessageSerializer());
     }
+
+    /**
+     * @return Enhanced {@link MessageFactory} allowing to create automated pre- and post- marshalling message serializer
+     * for {@link TcpDiscoveryMarshallableMessage}.
+     */
+    private MessageFactory enhanceMessageFactory(MessageFactory mf) {
+        if (cstDataMarshall == null || cstDataMarshallClsLdr == null)
+            return mf;
+
+        return new MessageFactory() {
+            @Override public void register(
+                short directType,
+                Supplier<Message> supplier,
+                MessageSerializer serializer
+            ) throws IgniteException {
+                if (supplier.get() instanceof TcpDiscoveryMarshallableMessage) {
+                    final MessageSerializer serializer0 = serializer;
+
+                    serializer = new MessageSerializer() {
+                        private Message curMarshallableMsg;
+
+                        @Override public boolean writeTo(Message msg, MessageWriter writer) {
+                            if (msg instanceof TcpDiscoveryMarshallableMessage && curMarshallableMsg == null) {
+                                curMarshallableMsg = msg;
+
+                                ((TcpDiscoveryMarshallableMessage)msg).prepareMarshal(cstDataMarshall);
+                            }
+
+                            boolean res = serializer0.writeTo(msg, writer);
+
+                            if (res && curMarshallableMsg != null) {
+                                assert msg instanceof TcpDiscoveryMarshallableMessage;
+
+                                curMarshallableMsg = null;
+                            }
+
+                            return res;
+                        }
+
+                        @Override public boolean readFrom(Message msg, MessageReader reader) {
+                            boolean res = serializer0.readFrom(msg, reader);
+
+                            if (res && msg instanceof TcpDiscoveryMarshallableMessage)
+                                ((TcpDiscoveryMarshallableMessage)msg).finishUnmarshal(cstDataMarshall, cstDataMarshallClsLdr);
+
+                            return res;
+                        }
+                    };
+                }
+
+                mf.register(directType, supplier, serializer);
+            }
+
+            @Override public void register(short directType, Supplier<Message> supplier) throws IgniteException {
+                mf.register(directType, supplier);
+            }
+
+            @Override public Message create(short type) {
+                return mf.create(type);
+            }
+
+            @Override public MessageSerializer serializer(short type) {
+                return mf.serializer(type);
+            }
+        };
+    }
 }

modules/core/src/main/java/org/apache/ignite/spi/discovery/tcp/ClientImpl.java

@@ -788,8 +788,6 @@ private static void sleepEx(long millis, Runnable before, Runnable after) throws
 
                     TcpDiscoveryJoinRequestMessage joinReqMsg = new TcpDiscoveryJoinRequestMessage(node, discoveryData);
 
-                    joinReqMsg.prepareMarshal(spi.marshaller());
-
                     TcpDiscoveryNode nodef = node;
 
                     joinReqMsg.spanContainer().span(
@@ -2241,7 +2239,7 @@ private void processNodeAddedMessage(TcpDiscoveryNodeAddedMessage msg) {
 
                         nodeAdded = true;
 
-                        if (msg.topologyHistory() != null)
+                        if (!F.isEmpty(msg.topologyHistory()))
                             topHist.putAll(msg.topologyHistory());
                     }
                     else {
@@ -2296,9 +2294,9 @@ private void processNodeAddFinishedMessage(TcpDiscoveryNodeAddFinishedMessage ms
                 }
             }
 
-            if (getLocalNodeId().equals(msg.nodeId())) {
+            if (getLocalNodeId().equals(msg.nodeId)) {
                 if (joining()) {
-                    DiscoveryDataPacket dataContainer = msg.clientDiscoData();
+                    DiscoveryDataPacket dataContainer = msg.clientDiscoData;
 
                     if (dataContainer != null)
                         spi.onExchange(dataContainer, U.resolveClassLoader(spi.ignite().configuration()));
@@ -2310,8 +2308,6 @@ private void processNodeAddFinishedMessage(TcpDiscoveryNodeAddFinishedMessage ms
                         delayDiscoData.clear();
                     }
 
-                    msg.finishUnmarshal(spi.marshaller(), U.resolveClassLoader(spi.ignite().configuration()));
-
                     locNode.setAttributes(msg.clientNodeAttributes());
 
                     clearNodeSensitiveData(locNode);
@@ -2355,7 +2351,7 @@ else if (log.isDebugEnabled())
             }
             else {
                 if (nodeAdded()) {
-                    TcpDiscoveryNode node = rmtNodes.get(msg.nodeId());
+                    TcpDiscoveryNode node = rmtNodes.get(msg.nodeId);
 
                     if (node == null) {
                         if (log.isDebugEnabled())
@@ -2548,6 +2544,10 @@ private void processClientReconnectMessage(TcpDiscoveryClientReconnectMessage ms
                 return;
 
             if (getLocalNodeId().equals(msg.creatorNodeId())) {
+                Collection<TcpDiscoveryAbstractMessage> pendingMsgs = msg.pendingMsgsMsg == null
+                    ? Collections.emptyList()
+                    : msg.pendingMsgsMsg.messages();
+
                 if (reconnector != null) {
                     assert msg.success() : msg;
 
@@ -2558,7 +2558,7 @@ private void processClientReconnectMessage(TcpDiscoveryClientReconnectMessage ms
 
                     reconnector = null;
 
-                    for (TcpDiscoveryAbstractMessage pendingMsg : msg.pendingMessages()) {
+                    for (TcpDiscoveryAbstractMessage pendingMsg : pendingMsgs) {
                         if (log.isDebugEnabled())
                             log.debug("Process pending message on reconnect [msg=" + pendingMsg + ']');
 
@@ -2568,7 +2568,7 @@ private void processClientReconnectMessage(TcpDiscoveryClientReconnectMessage ms
                 else {
                     if (joinLatch.getCount() > 0) {
                         if (msg.success()) {
-                            for (TcpDiscoveryAbstractMessage pendingMsg : msg.pendingMessages()) {
+                            for (TcpDiscoveryAbstractMessage pendingMsg : pendingMsgs) {
                                 if (log.isDebugEnabled())
                                     log.debug("Process pending message on connect [msg=" + pendingMsg + ']');