IGNITE-27722 Split TcpDiscoveryIoSession to fix race in RingMessageWorker (#12731)

(cherry picked from commit a182326)

Author: sergey-chugunov-1985

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

@@ -100,6 +100,7 @@
 import org.apache.ignite.internal.util.typedef.C1;
 import org.apache.ignite.internal.util.typedef.F;
 import org.apache.ignite.internal.util.typedef.P1;
+import org.apache.ignite.internal.util.typedef.T2;
 import org.apache.ignite.internal.util.typedef.X;
 import org.apache.ignite.internal.util.typedef.internal.LT;
 import org.apache.ignite.internal.util.typedef.internal.S;
@@ -2834,6 +2835,16 @@ private class RingMessageWorker extends MessageWorker<TcpDiscoveryAbstractMessag
         /** Socket. */
         private Socket sock;
 
+        // This serializer is used exclusively for serializing messages sent to clients,
+        // as it represents a special case within the RingMessageWorker workflow.
+        // Generally, both serialization and deserialization of messages should be handled by TcpDiscoveryIoSession.
+        // However, there are scenarios where the session is not available, yet messages still need to be sent to clients.
+        // A typical example is a single server with one or more connected clients.
+        // To address this, we use TcpDiscoveryMessageSerializer, which includes some code copied from TcpDiscoveryIoSession
+        // and can be instantiated independently of any active session.
+        /** */
+        private final TcpDiscoveryMessageSerializer clientMsgSer = new TcpDiscoveryMessageSerializer(spi);
+
         /** IO session. */
         private TcpDiscoveryIoSession ses;
 
@@ -3241,19 +3252,36 @@ private void sendMessageToClients(TcpDiscoveryAbstractMessage msg) {
                 if (spi.ensured(msg))
                     msgHist.add(msg);
 
+                if (clientMsgWorkers.isEmpty())
+                    return;
+
+                byte[] msgBytes = null;
+
+                if (!(msg instanceof TcpDiscoveryNodeAddedMessage)) {
+                    try {
+                        msgBytes = clientMsgSer.serializeMessage(msg);
+                    }
+                    catch (IgniteCheckedException | IOException e) {
+                        U.error(log, "Failed to serialize message: " + msg, e);
+
+                        return;
+                    }
+                }
+
                 for (ClientMessageWorker clientMsgWorker : clientMsgWorkers.values()) {
+                    TcpDiscoveryAbstractMessage msg0 = msg;
+
                     if (msg instanceof TcpDiscoveryNodeAddedMessage) {
                         TcpDiscoveryNodeAddedMessage nodeAddedMsg = (TcpDiscoveryNodeAddedMessage)msg;
 
                         if (clientMsgWorker.clientNodeId.equals(nodeAddedMsg.node().id())) {
-                            msg = new TcpDiscoveryNodeAddedMessage(nodeAddedMsg);
+                            msg0 = new TcpDiscoveryNodeAddedMessage(nodeAddedMsg);
 
-                            prepareNodeAddedMessage(msg, clientMsgWorker.clientNodeId, null);
+                            prepareNodeAddedMessage(msg0, clientMsgWorker.clientNodeId, null);
                         }
                     }
 
-                    // TODO Investigate possible optimizations: https://issues.apache.org/jira/browse/IGNITE-27722
-                    clientMsgWorker.addMessage(msg);
+                    clientMsgWorker.addMessage(msg0, msgBytes);
                 }
             }
         }
@@ -7584,12 +7612,20 @@ private class StatisticsPrinter extends IgniteSpiThread {
     }
 
     /** */
-    private class ClientMessageWorker extends MessageWorker<TcpDiscoveryAbstractMessage> {
+    private class ClientMessageWorker extends MessageWorker<T2<TcpDiscoveryAbstractMessage, byte[]>> {
         /** Node ID. */
         private final UUID clientNodeId;
 
+        // The code responsible for sending and receiving messages to and from client nodes represents a special case in ServerImpl,
+        // as it is split into two separate components.
+        // One part, ClientMessageWorker, handles only message sending to clients and does not process responses.
+        // The other part, which reads messages from clients, is implemented in SocketReader.
+        // Due to this separation, we don't require a full TcpDiscoveryIoSession here
+        // and can instead extract just the message-writing functionality.
+        // At the same time, we aim to keep both reading and writing logic encapsulated within TcpDiscoveryIoSession.
+        // As a result, we need to copy some code from TcpDiscoveryIoSession into the new class, TcpDiscoveryMessageSerializer.
         /** */
-        private final TcpDiscoveryIoSession ses;
+        private final TcpDiscoveryMessageSerializer clientMsgSer;
 
         /** Socket. */
         private final Socket sock;
@@ -7619,7 +7655,7 @@ private ClientMessageWorker(Socket sock, UUID clientNodeId, IgniteLogger log) {
             this.sock = sock;
             this.clientNodeId = clientNodeId;
 
-            ses = createSession(sock);
+            clientMsgSer = new TcpDiscoveryMessageSerializer(spi);
 
             lastMetricsUpdateMsgTimeNanos = System.nanoTime();
         }
@@ -7651,10 +7687,20 @@ void metrics(ClusterMetrics metrics) {
          * @param msg Message.
          */
         void addMessage(TcpDiscoveryAbstractMessage msg) {
+            addMessage(msg, null);
+        }
+
+        /**
+         * @param msg Message.
+         * @param msgBytes Optional message bytes.
+         */
+        void addMessage(TcpDiscoveryAbstractMessage msg, @Nullable byte[] msgBytes) {
+            T2<TcpDiscoveryAbstractMessage, byte[]> t = new T2<>(msg, msgBytes);
+
             if (msg.highPriority())
-                queue.addFirst(msg);
+                queue.addFirst(t);
             else
-                queue.add(msg);
+                queue.add(t);
 
             DebugLogger log = messageLogger(msg);
 
@@ -7663,14 +7709,14 @@ void addMessage(TcpDiscoveryAbstractMessage msg) {
         }
 
         /** {@inheritDoc} */
-        @Override protected void processMessage(TcpDiscoveryAbstractMessage msg) {
+        @Override protected void processMessage(T2<TcpDiscoveryAbstractMessage, byte[]> msgT) {
             boolean success = false;
 
+            TcpDiscoveryAbstractMessage msg = msgT.get1();
+
             try {
                 assert msg.verified() : msg;
 
-                byte[] msgBytes = ses.serializeMessage(msg);
-
                 DebugLogger msgLog = messageLogger(msg);
 
                 if (msg instanceof TcpDiscoveryClientAckResponse) {
@@ -7692,8 +7738,8 @@ else if (msgLog.isDebugEnabled()) {
                                 + getLocalNodeId() + ", rmtNodeId=" + clientNodeId + ", msg=" + msg + ']');
                         }
 
-                        spi.writeToSocket(sock, msg, msgBytes, spi.failureDetectionTimeoutEnabled() ?
-                            spi.clientFailureDetectionTimeout() : spi.getSocketTimeout());
+                        writeToSocket(msgT, spi.failureDetectionTimeoutEnabled() ? spi.clientFailureDetectionTimeout() :
+                            spi.getSocketTimeout());
                     }
                 }
                 else {
@@ -7704,7 +7750,7 @@ else if (msgLog.isDebugEnabled()) {
 
                     assert topologyInitialized(msg) : msg;
 
-                    spi.writeToSocket(sock, msg, msgBytes, spi.getEffectiveSocketTimeout(false));
+                    writeToSocket(msgT, spi.getEffectiveSocketTimeout(false));
                 }
 
                 boolean clientFailed = msg instanceof TcpDiscoveryNodeFailedMessage &&
@@ -7733,6 +7779,17 @@ else if (msgLog.isDebugEnabled()) {
             }
         }
 
+        /**
+         * @param msgT Message tuple.
+         * @param timeout Timeout.
+         */
+        private void writeToSocket(T2<TcpDiscoveryAbstractMessage, byte[]> msgT, long timeout)
+            throws IgniteCheckedException, IOException {
+            byte[] msgBytes = msgT.get2() == null ? clientMsgSer.serializeMessage(msgT.get1()) : msgT.get2();
+
+            spi.writeToSocket(sock, msgT.get1(), msgBytes, timeout);
+        }
+
         /**
          * @param msg Message.
          * @return {@code True} if topology initialized.

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

@@ -20,7 +20,6 @@
 import java.io.BufferedInputStream;
 import java.io.BufferedOutputStream;
 import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
 import java.io.EOFException;
 import java.io.IOException;
 import java.io.InputStream;
@@ -71,7 +70,7 @@ public class TcpDiscoveryIoSession {
     static final byte MESSAGE_SERIALIZATION = (byte)2;
 
     /** */
-    private final TcpDiscoverySpi spi;
+    final TcpDiscoverySpi spi;
 
     /** Loads discovery messages classes during java deserialization. */
     private final ClassLoader clsLdr;
@@ -80,7 +79,7 @@ public class TcpDiscoveryIoSession {
     private final Socket sock;
 
     /** */
-    private final DirectMessageWriter msgWriter;
+    final DirectMessageWriter msgWriter;
 
     /** */
     private final DirectMessageReader msgReader;
@@ -92,7 +91,7 @@ public class TcpDiscoveryIoSession {
     private final CompositeInputStream in;
 
     /** Intermediate buffer for serializing discovery messages. */
-    private final ByteBuffer msgBuf;
+    final ByteBuffer msgBuf;
 
     /**
      * Creates a new discovery I/O session bound to the given socket.
@@ -238,25 +237,6 @@ <T> T readMessage() throws IgniteCheckedException, IOException {
         }
     }
 
-    /**
-     * Serializes a discovery message into a byte array.
-     *
-     * @param msg Discovery message to serialize.
-     * @return Serialized byte array containing the message data.
-     * @throws IgniteCheckedException If serialization fails.
-     * @throws IOException If serialization fails.
-     */
-    byte[] serializeMessage(TcpDiscoveryAbstractMessage msg) throws IgniteCheckedException, IOException {
-        if (!(msg instanceof Message))
-            return U.marshal(spi.marshaller(), msg);
-
-        try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
-            serializeMessage((Message)msg, out);
-
-            return out.toByteArray();
-        }
-    }
-
     /** @return Socket. */
     public Socket socket() {
         return sock;
@@ -269,7 +249,7 @@ public Socket socket() {
      * @param out Output stream to write serialized message.
      * @throws IOException If serialization fails.
      */
-    private void serializeMessage(Message m, OutputStream out) throws IOException {
+    void serializeMessage(Message m, OutputStream out) throws IOException {
         MessageSerializer msgSer = spi.messageFactory().serializer(m.directType());
 
         msgWriter.reset();

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

@@ -0,0 +1,71 @@
+/*
+ * 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.spi.discovery.tcp;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.net.Socket;
+import org.apache.ignite.IgniteCheckedException;
+import org.apache.ignite.internal.util.typedef.internal.U;
+import org.apache.ignite.plugin.extensions.communication.Message;
+import org.apache.ignite.plugin.extensions.communication.MessageSerializer;
+import org.apache.ignite.spi.discovery.tcp.messages.TcpDiscoveryAbstractMessage;
+
+/**
+ * Class is responsible for serializing discovery messages using RU-ready {@link MessageSerializer} mechanism.
+ * <p>
+ * It is used in a special case: when server wants to send discovery messages to clients, it may not have a {@link TcpDiscoveryIoSession}
+ * to serialize the messages.
+ * This class enables server to serialize discovery messages anyway, duplicating serialization code from {@link TcpDiscoveryIoSession}.
+ */
+class TcpDiscoveryMessageSerializer extends TcpDiscoveryIoSession {
+    /**
+     * @param spi Discovery SPI instance.
+     */
+    public TcpDiscoveryMessageSerializer(TcpDiscoverySpi spi) {
+        super(new Socket() {
+            @Override public OutputStream getOutputStream() throws IOException {
+                return null;
+            }
+
+            @Override public InputStream getInputStream() throws IOException {
+                return null;
+            }
+        }, spi);
+    }
+
+    /**
+     * Serializes a discovery message into a byte array.
+     *
+     * @param msg Discovery message to serialize.
+     * @return Serialized byte array containing the message data.
+     * @throws IgniteCheckedException If serialization fails.
+     * @throws IOException If serialization fails.
+     */
+    byte[] serializeMessage(TcpDiscoveryAbstractMessage msg) throws IgniteCheckedException, IOException {
+        if (!(msg instanceof Message))
+            return U.marshal(spi.marshaller(), msg);
+
+        try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+            serializeMessage((Message)msg, out);
+
+            return out.toByteArray();
+        }
+    }
+}