IGNITE-28784 : Send operation context attributes via Communication

modules/commons/src/main/java/org/apache/ignite/internal/thread/context/OperationContext.java

@@ -322,7 +322,7 @@ private static class AttributeValueHolder<T> {
     }
 
     /** Allows to change multiple attribute values in a single update operation and skip updates that changes nothing. */
-    private static class ContextUpdater {
+    static class ContextUpdater {
         /** */
         private static final int INIT_UPDATES_CAPACITY = 3;
 

modules/core/src/main/java/org/apache/ignite/internal/CoreMessagesProvider.java

@@ -666,6 +666,10 @@ public CoreMessagesProvider(Marshaller dfltMarsh, Marshaller schemaAwareMarsh, C
         withNoSchema(PartitionHashRecord.class);
         withNoSchema(TransactionsHashRecord.class);
 
+        // [13400 - 13500]: Operation context messages.
+        msgIdx = 13400;
+        withNoSchema(DistributedOperationContextMessage.class);
+
         assert msgIdx <= MAX_MESSAGE_ID;
     }
 

modules/core/src/main/java/org/apache/ignite/internal/DistributedOperationContextMessage.java

@@ -0,0 +1,38 @@
+/*
+ * 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.internal;
+
+import java.util.List;
+import org.apache.ignite.internal.thread.context.OperationContext;
+import org.apache.ignite.plugin.extensions.communication.Message;
+
+/** Transport for {@link OperationContext} distributed attributes. */
+public class DistributedOperationContextMessage implements Message {
+    /** Values of operation context attributes. */
+    @Order(0)
+    public List<Message> vals;
+
+    /** Bitmask of effective attributes ids. */
+    @Order(1)
+    public byte idBitmap;
+
+    /** Empty constructor for serialization purposes. */
+    public DistributedOperationContextMessage() {
+        // No-op.
+    }
+}

modules/core/src/main/java/org/apache/ignite/internal/managers/communication/GridIoManager.java

@@ -102,6 +102,7 @@
 import org.apache.ignite.internal.processors.tracing.MTC.TraceSurroundings;
 import org.apache.ignite.internal.processors.tracing.Span;
 import org.apache.ignite.internal.processors.tracing.SpanTags;
+import org.apache.ignite.internal.thread.context.DistributedOperationContextManager;
 import org.apache.ignite.internal.thread.context.Scope;
 import org.apache.ignite.internal.util.GridBoundedConcurrentLinkedHashSet;
 import org.apache.ignite.internal.util.IgniteUtils;
@@ -455,10 +456,14 @@ public void resetMetrics() {
 
         ioMetric.register(RCVD_BYTES_CNT, spi::getReceivedBytesCount, "Received bytes count.");
 
-        getSpi().setListener(commLsnr = new CommunicationListenerEx<Object>() {
+        getSpi().setListener(commLsnr = new CommunicationListenerEx<>() {
             @Override public void onMessage(UUID nodeId, Object msg, IgniteRunnable msgC) {
                 try {
-                    onMessage0(nodeId, (GridIoMessage)msg, msgC);
+                    GridIoMessage msg0 = (GridIoMessage)msg;
+
+                    try (Scope ignored = DistributedOperationContextManager.instance().restoreDistributedAttributes(msg0.opCtxMsg)) {
+                        onMessage0(nodeId, msg0, msgC);
+                    }
                 }
                 catch (ClassCastException ignored) {
                     U.error(log, "Communication manager received message of unknown type (will ignore): " +
@@ -2037,16 +2042,22 @@ private long getInverseConnectionWaitTimeout() {
         long timeout,
         boolean skipOnTimeout
     ) {
+        GridIoMessage res;
+
         if (ctx.security().enabled()) {
             UUID secSubjId = null;
 
             if (!ctx.security().isDefaultContext())
                 secSubjId = ctx.security().securityContext().subject().id();
 
-            return new GridIoSecurityAwareMessage(secSubjId, plc, topic, msg, ordered, timeout, skipOnTimeout);
+            res = new GridIoSecurityAwareMessage(secSubjId, plc, topic, msg, ordered, timeout, skipOnTimeout);
         }
+        else
+            res = new GridIoMessage(plc, topic, msg, ordered, timeout, skipOnTimeout);
+
+        res.opCtxMsg = DistributedOperationContextManager.instance().collectDistributedAttributes();
 
-        return new GridIoMessage(plc, topic, msg, ordered, timeout, skipOnTimeout);
+        return res;
     }
 
     /**

modules/core/src/main/java/org/apache/ignite/internal/managers/communication/GridIoMessage.java

@@ -17,6 +17,7 @@
 
 package org.apache.ignite.internal.managers.communication;
 
+import org.apache.ignite.internal.DistributedOperationContextMessage;
 import org.apache.ignite.internal.ExecutorAwareMessage;
 import org.apache.ignite.internal.GridTopicMessage;
 import org.apache.ignite.internal.Order;
@@ -64,6 +65,11 @@ public class GridIoMessage implements Message, SpanTransport {
     @Order(6)
     byte[] span;
 
+    /** Effective operation context attributes. */
+    @Order(7)
+    @GridToStringInclude
+    public @Nullable DistributedOperationContextMessage opCtxMsg;
+
     /**
      * Default constructor.
      */

modules/core/src/main/java/org/apache/ignite/internal/thread/context/DistributedOperationContextManager.java

@@ -0,0 +1,116 @@
+/*
+ * 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.internal.thread.context;
+
+import java.util.ArrayList;
+import java.util.Map;
+import java.util.concurrent.ConcurrentSkipListMap;
+import org.apache.ignite.IgniteException;
+import org.apache.ignite.internal.DistributedOperationContextMessage;
+import org.apache.ignite.internal.util.typedef.F;
+import org.apache.ignite.plugin.extensions.communication.Message;
+import org.jetbrains.annotations.Nullable;
+
+/** */
+public class DistributedOperationContextManager {
+    /** */
+    static final byte MAX_DISTRIBUTED_ATTR_CNT = 7;
+
+    /** */
+    private static final DistributedOperationContextManager INSTANCE = new DistributedOperationContextManager();
+
+    /** Attributes by their id. */
+    private final Map<Byte, OperationContextAttribute<? extends Message>> attrs = new ConcurrentSkipListMap<>();
+
+    /** */
+    public static DistributedOperationContextManager instance() {
+        return INSTANCE;
+    }
+
+    /** */
+    public <T extends Message> OperationContextAttribute<T> createDistributedAttriubte(byte id, @Nullable T initVal) {
+        assert id >= 0;
+
+        if (attrs.size() == OperationContextAttribute.MAX_ATTR_CNT) {
+            throw new IgniteException("Maximum number of distributed attributes is exceeded [max="
+                + OperationContextAttribute.MAX_ATTR_CNT + "].");
+        }
+
+        return (OperationContextAttribute<T>)attrs.compute(id, (id0, attr0) -> {
+            if (attr0 != null)
+                throw new IgniteException("Duplicated distributed attribute id [id=" + id + "].");
+
+            return OperationContextAttribute.newInstance(initVal);
+        });
+    }
+
+    /** */
+    public @Nullable DistributedOperationContextMessage collectDistributedAttributes() {
+        DistributedOperationContextMessage res = null;
+
+        for (Map.Entry<Byte, OperationContextAttribute<? extends Message>> e : attrs.entrySet()) {
+            OperationContextAttribute<? extends Message> attr = e.getValue();
+
+            Message curVal = OperationContext.get(attr);
+
+            assert attr.initialValue() == null || curVal == null || curVal.getClass().isAssignableFrom(attr.initialValue().getClass());
+
+            if (curVal != attr.initialValue()) {
+                if (res == null) {
+                    res = new DistributedOperationContextMessage();
+
+                    res.vals = new ArrayList<>(MAX_DISTRIBUTED_ATTR_CNT / 2);
+                }
+
+                byte mask = (byte)(1 << e.getKey());
+
+                assert (res.idBitmap & mask) == 0;
+
+                res.vals.add(curVal);
+                res.idBitmap |= mask;
+            }
+        }
+
+        return res;
+    }
+
+    /** */
+    public Scope restoreDistributedAttributes(@Nullable DistributedOperationContextMessage msg) {
+        if (msg == null)
+            return Scope.NOOP_SCOPE;
+
+        assert msg.idBitmap != 0;
+        assert !F.isEmpty(msg.vals);
+        assert msg.vals.size() <= MAX_DISTRIBUTED_ATTR_CNT;
+
+        OperationContext.ContextUpdater updater = OperationContext.ContextUpdater.create();
+
+        for (byte valIdx = 0, maskIdx = 0; valIdx < msg.vals.size(); ++valIdx) {
+            Message curVal = msg.vals.get(valIdx);
+
+            while ((msg.idBitmap & (1 << maskIdx)) == 0) {
+                assert maskIdx <= MAX_DISTRIBUTED_ATTR_CNT;
+
+                ++maskIdx;
+            }
+
+            updater.set((OperationContextAttribute<Message>)attrs.get(maskIdx++), curVal);
+        }
+
+        return updater.apply();
+    }
+}

modules/core/src/main/java/org/apache/ignite/spi/communication/tcp/TcpCommunicationSpi.java

@@ -616,7 +616,7 @@ private void dumpInfo(StringBuilder sb, UUID dstNodeId) {
             ctxInitLatch,
             client,
             igniteExSupplier,
-            new CommunicationListener<Message>() {
+            new CommunicationListener<>() {
                 @Override public void onMessage(UUID nodeId, Message msg, IgniteRunnable msgC) {
                     notifyListener(nodeId, msg, msgC);
                 }
@@ -651,7 +651,7 @@ private void dumpInfo(StringBuilder sb, UUID dstNodeId) {
             getWorkersRegistry(ignite),
             ignite instanceof IgniteEx ? ((IgniteEx)ignite).context().metric() : null,
             this::createTcpClient,
-            new CommunicationListenerEx<Message>() {
+            new CommunicationListenerEx<>() {
                 @Override public void onMessage(UUID nodeId, Message msg, IgniteRunnable msgC) {
                     notifyListener(nodeId, msg, msgC);
                 }

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

@@ -70,6 +70,8 @@
 import org.apache.ignite.internal.processors.tracing.messages.SpanContainer;
 import org.apache.ignite.internal.processors.tracing.messages.TraceableMessage;
 import org.apache.ignite.internal.processors.tracing.messages.TraceableMessagesTable;
+import org.apache.ignite.internal.thread.context.DistributedOperationContextManager;
+import org.apache.ignite.internal.thread.context.Scope;
 import org.apache.ignite.internal.util.future.GridFutureAdapter;
 import org.apache.ignite.internal.util.tostring.GridToStringExclude;
 import org.apache.ignite.internal.util.typedef.F;
@@ -1310,6 +1312,8 @@ private class SocketWriter extends IgniteSpiThread {
          * @param msg Message.
          */
         private void sendMessage(TcpDiscoveryAbstractMessage msg) {
+            msg.opCtxMsg = DistributedOperationContextManager.instance().collectDistributedAttributes();
+
             synchronized (mux) {
                 queue.add(msg);
 
@@ -2001,7 +2005,11 @@ else if (discoMsg instanceof TcpDiscoveryCheckFailedMessage)
                             }
                         }
 
-                        processDiscoveryMessage((TcpDiscoveryAbstractMessage)msg);
+                        TcpDiscoveryAbstractMessage msg0 = (TcpDiscoveryAbstractMessage)msg;
+
+                        try (Scope ignored = DistributedOperationContextManager.instance().restoreDistributedAttributes(msg0.opCtxMsg)) {
+                            processDiscoveryMessage(msg0);
+                        }
                     }
                 }
             }

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

@@ -95,6 +95,8 @@
 import org.apache.ignite.internal.processors.tracing.messages.SpanContainer;
 import org.apache.ignite.internal.processors.tracing.messages.TraceableMessage;
 import org.apache.ignite.internal.processors.tracing.messages.TraceableMessagesTable;
+import org.apache.ignite.internal.thread.context.DistributedOperationContextManager;
+import org.apache.ignite.internal.thread.context.Scope;
 import org.apache.ignite.internal.thread.pool.IgniteThreadPoolExecutor;
 import org.apache.ignite.internal.util.GridBoundedLinkedHashSet;
 import org.apache.ignite.internal.util.GridConcurrentHashSet;
@@ -3046,8 +3048,10 @@ void addMessage(TcpDiscoveryAbstractMessage msg, boolean ignoreHighPriority, boo
                 return;
             }
 
-            if (msg instanceof TraceableMessage) {
-                TraceableMessage tMsg = (TraceableMessage)msg;
+            if (!fromSocket)
+                msg.opCtxMsg = DistributedOperationContextManager.instance().collectDistributedAttributes();
+
+            if (msg instanceof TraceableMessage tMsg) {
 
                 // If we read this message from socket.
                 if (fromSocket)
@@ -3173,11 +3177,8 @@ protected void runTasks() {
                 task.run();
         }
 
-        /** {@inheritDoc} */
-        @Override protected void processMessage(TcpDiscoveryAbstractMessage msg) {
-            if (msg == WAKEUP)
-                return;
-
+        /** */
+        private void processMessage0(TcpDiscoveryAbstractMessage msg) {
             notifiedDiscovery.set(false);
 
             if (msg instanceof TraceableMessage) {
@@ -3315,6 +3316,16 @@ else if (msg instanceof TcpDiscoveryAuthFailedMessage)
             }
         }
 
+        /** {@inheritDoc} */
+        @Override protected void processMessage(TcpDiscoveryAbstractMessage msg) {
+            if (msg == WAKEUP)
+                return;
+
+            try (Scope ignored = DistributedOperationContextManager.instance().restoreDistributedAttributes(msg.opCtxMsg)) {
+                processMessage0(msg);
+            }
+        }
+
         /**
          * Processes authentication failed message.
          *

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

@@ -52,7 +52,6 @@ public int port() {
         return port;
     }
 
-
     /** {@inheritDoc} */
     @Override public String toString() {
         return S.toString(InetSocketAddressMessage.class, this);

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

@@ -21,6 +21,7 @@
 import java.util.HashSet;
 import java.util.Set;
 import java.util.UUID;
+import org.apache.ignite.internal.DistributedOperationContextMessage;
 import org.apache.ignite.internal.Order;
 import org.apache.ignite.internal.util.tostring.GridToStringExclude;
 import org.apache.ignite.internal.util.tostring.GridToStringInclude;
@@ -76,6 +77,11 @@ public abstract class TcpDiscoveryAbstractMessage implements Message {
     @Order(4)
     Set<UUID> failedNodes;
 
+    /** Operation context attributes message. */
+    @GridToStringInclude
+    @Order(5)
+    public @Nullable DistributedOperationContextMessage opCtxMsg;
+
     /**
      * Default no-arg constructor for {@link Externalizable} interface.
      */