IGNITE-28902 Fixed Operation Context propagation for discover ack messages via TCP Discovery SPI

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

@@ -259,6 +259,7 @@
 import org.apache.ignite.internal.processors.service.ServiceSingleNodeDeploymentResultBatch;
 import org.apache.ignite.internal.processors.service.ServiceTopology;
 import org.apache.ignite.internal.processors.service.ServiceUndeploymentRequest;
+import org.apache.ignite.internal.thread.context.OperationContextMessage;
 import org.apache.ignite.internal.util.GridByteArrayList;
 import org.apache.ignite.internal.util.GridIntList;
 import org.apache.ignite.internal.util.GridPartitionStateMap;

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

@@ -459,7 +459,7 @@ public void resetMetrics() {
                 try {
                     GridIoMessage msg0 = (GridIoMessage)msg;
 
-                    try (Scope ignored = ctx.operationContextDispatcher().restoreDistributedAttributes(msg0.opCtxMsg)) {
+                    try (Scope ignored = ctx.operationContextDispatcher().restoreRemoteAttributeValues(msg0.opCtxMsg)) {
                         onMessage0(nodeId, msg0, msgC);
                     }
                 }
@@ -2051,7 +2051,7 @@ public GridIoMessage createGridIoMessage(
 
         res = new GridIoMessage(plc, topic, msg, ordered, timeout, skipOnTimeout);
 
-        res.opCtxMsg = ctx.operationContextDispatcher().collectDistributedAttributes();
+        res.opCtxMsg = ctx.operationContextDispatcher().collectDistributedAttributeValues();
 
         return res;
     }

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

@@ -19,11 +19,11 @@
 
 import org.apache.ignite.internal.ExecutorAwareMessage;
 import org.apache.ignite.internal.GridTopicMessage;
-import org.apache.ignite.internal.OperationContextMessage;
 import org.apache.ignite.internal.Order;
 import org.apache.ignite.internal.processors.cache.GridCacheMessage;
 import org.apache.ignite.internal.processors.datastreamer.DataStreamerRequest;
 import org.apache.ignite.internal.processors.tracing.messages.SpanTransport;
+import org.apache.ignite.internal.thread.context.OperationContextMessage;
 import org.apache.ignite.internal.util.nio.GridNioServer.MessageWrapper;
 import org.apache.ignite.internal.util.tostring.GridToStringInclude;
 import org.apache.ignite.internal.util.typedef.internal.S;

modules/core/src/main/java/org/apache/ignite/internal/processors/security/IgniteSecurityProcessor.java

@@ -56,7 +56,7 @@
 import static org.apache.ignite.internal.processors.security.SecurityUtils.MSG_SEC_PROC_CLS_IS_INVALID;
 import static org.apache.ignite.internal.processors.security.SecurityUtils.hasSecurityManager;
 import static org.apache.ignite.internal.processors.security.SecurityUtils.nodeSecurityContext;
-import static org.apache.ignite.internal.thread.context.DistributedOperationContextAttribute.SECURITY;
+import static org.apache.ignite.internal.thread.context.DistributedAttributeRegistry.SECURITY;
 import static org.apache.ignite.plugin.security.SecurityPermission.ADMIN_USER_ACCESS;
 import static org.apache.ignite.plugin.security.SecurityPermission.JOIN_AS_SERVER;
 
@@ -253,7 +253,7 @@ private SecurityContext securityContext(UUID subjId) {
     @Override public void start() throws IgniteCheckedException {
         super.start();
 
-        ctx.operationContextDispatcher().registerDistributedAttribute(SECURITY.id(), SEC_CTX_ATTR);
+        ctx.operationContextDispatcher().registerDistributedAttribute(SECURITY, SEC_CTX_ATTR);
 
         ctx.addNodeAttribute(ATTR_GRID_SEC_PROC_CLASS, secPrc.getClass().getName());
 

modules/core/src/main/java/org/apache/ignite/internal/processors/security/SecurityContextWrapper.java

@@ -19,16 +19,16 @@
 
 import java.util.UUID;
 import org.apache.ignite.internal.Order;
-import org.apache.ignite.internal.thread.context.DistributedOperationContextAttribute;
+import org.apache.ignite.internal.thread.context.DistributedAttributeRegistry;
 import org.apache.ignite.internal.thread.context.OperationContextDispatcher;
 import org.apache.ignite.plugin.extensions.communication.Message;
 import org.apache.ignite.plugin.security.SecuritySubject;
 
 /**
  * {@link SecurityContext} attribute value holder and message for {@link SecuritySubject}'s id.
  *
- * @see OperationContextDispatcher#collectDistributedAttributes()
- * @see DistributedOperationContextAttribute#SECURITY
+ * @see OperationContextDispatcher#collectDistributedAttributeValues()
+ * @see DistributedAttributeRegistry#SECURITY
  */
 public class SecurityContextWrapper implements Message {
     /** A value of {@link SecuritySubject#id()} */

modules/core/src/main/java/org/apache/ignite/internal/thread/context/DistributedOperationContextAttribute.java → modules/core/src/main/java/org/apache/ignite/internal/thread/context/DistributedAttributeRegistry.java

@@ -18,19 +18,12 @@
 package org.apache.ignite.internal.thread.context;
 
 import org.apache.ignite.internal.processors.security.SecurityContext;
-import org.apache.ignite.internal.processors.security.SecurityContextWrapper;
 
-/** Ids of Ignite's known distributed operation context attributes. */
-public enum DistributedOperationContextAttribute {
-    /**
-     * Distributed {@link SecurityContext}.
-     *
-     * @see SecurityContextWrapper
-     */
-    SECURITY;
-
-    /** Cluster-wide id of distributed attribute. */
-    public byte id() {
-        return (byte)ordinal();
-    }
+/**
+ * Declares reserved distributed IDs used to consistently identify {@link OperationContext} attributes across
+ * all nodes in the cluster.
+ */
+public class DistributedAttributeRegistry {
+    /** Reserved for {@link SecurityContext} propagation. */
+    public static final byte SECURITY = 0;
 }

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

@@ -17,11 +17,9 @@
 package org.apache.ignite.internal.thread.context;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ConcurrentSkipListMap;
 import org.apache.ignite.IgniteException;
-import org.apache.ignite.internal.OperationContextMessage;
 import org.apache.ignite.internal.util.typedef.F;
 import org.apache.ignite.plugin.extensions.communication.Message;
 import org.jetbrains.annotations.Nullable;
@@ -44,17 +42,16 @@
  *
  * @see OperationContext
  * @see OperationContextMessage
- * @see DistributedOperationContextAttribute
  */
 public class OperationContextDispatcher {
     /** Maximal number of supported distributed attributes. */
     static final byte MAX_ATTRS_CNT = Byte.SIZE;
 
     /** Registered distributed attributes by their cluster-wide id. */
-    private final Map<Byte, OperationContextAttribute<? extends Message>> attrs = new ConcurrentSkipListMap<>();
+    private volatile OperationContextAttribute<? extends Message>[] registeredAttrs = new OperationContextAttribute[0];
 
     /** Whether the registration of new distributed attributes is allowed. */
-    private volatile boolean regFinished;
+    private boolean regFinished;
 
     /**
      * Registers an attribute of {@link OperationContext} with the specified distributed ID.
@@ -64,71 +61,90 @@ public class OperationContextDispatcher {
      *
      * <p>Registered attribute value is automatically captured and propagated between cluster nodes
      * during the messages transmission.</p>
+     *
+     * @see DistributedAttributeRegistry
      */
-    public <T extends Message> void registerDistributedAttribute(int id, OperationContextAttribute<T> attr) {
+    public synchronized <T extends Message> void registerDistributedAttribute(int id, OperationContextAttribute<T> attr) {
         if (regFinished)
             throw new IgniteException("Initialization of distributed operation context attributes has already finished.");
 
-        assert id >= 0 && id < MAX_ATTRS_CNT : "Invalid distributed attributed id [id=" + id + ']';
+        assert 0 <= id && id < MAX_ATTRS_CNT : "Invalid distributed attributed id [id=" + id + ']';
+
+        OperationContextAttribute<? extends Message>[] locRegisteredAttrs = registeredAttrs;
+
+        OperationContextAttribute<? extends Message>[] copy = Arrays.copyOf(
+            locRegisteredAttrs,
+            Math.max(locRegisteredAttrs.length, id + 1));
 
-        if (attrs.putIfAbsent((byte)id, attr) != null)
+        if (copy[id] != null)
             throw new IgniteException("Duplicated distributed attribute id [id=" + id + ']');
+
+        copy[id] = attr;
+
+        registeredAttrs = copy;
     }
 
     /**
      * Collects the values of all distributed {@link OperationContextAttribute}s registered by this dispatcher.
      *
      * @see OperationContext#get(OperationContextAttribute)
      */
-    public @Nullable OperationContextMessage collectDistributedAttributes() {
-        OperationContextMessage res = null;
+    public @Nullable OperationContextMessage collectDistributedAttributeValues() {
+        OperationContextAttribute<? extends Message>[] locRegisteredAttrs = registeredAttrs;
+
+        if (locRegisteredAttrs.length == 0)
+            return null;
+
+        byte bitmap = 0;
         List<Message> vals = null;
 
-        for (Map.Entry<Byte, OperationContextAttribute<? extends Message>> e : attrs.entrySet()) {
-            OperationContextAttribute<? extends Message> attr = e.getValue();
+        for (int id = 0; id < locRegisteredAttrs.length; id++) {
+            OperationContextAttribute<? extends Message> attr = locRegisteredAttrs[id];
+
+            if (attr == null)
+                continue;
 
             Message curVal = OperationContext.get(attr);
 
-            if (curVal != attr.initialValue()) {
-                if (res == null) {
-                    res = new OperationContextMessage();
+            if (curVal == attr.initialValue())
+                continue;
 
-                    vals = new ArrayList<>(MAX_ATTRS_CNT / 2);
-                }
+            if (vals == null)
+                vals = new ArrayList<>(MAX_ATTRS_CNT / 2);
 
-                byte mask = (byte)(1 << e.getKey());
+            byte mask = (byte)(1 << id);
 
-                assert (res.idBitmap & mask) == 0;
+            assert (bitmap & mask) == 0;
 
-                vals.add(curVal);
-                res.idBitmap |= mask;
-            }
+            vals.add(curVal);
+            bitmap |= mask;
         }
 
-        if (res != null)
-            res.vals = vals.toArray(new Message[vals.size()]);
-
-        return res;
+        return bitmap == 0 ? null : new OperationContextMessage(bitmap, vals.toArray(Message[]::new));
     }
 
     /** Restores distributed {@link OperationContextAttribute} values received from a remote node. */
-    public Scope restoreDistributedAttributes(@Nullable OperationContextMessage msg) {
+    public Scope restoreRemoteAttributeValues(@Nullable OperationContextMessage msg) {
         if (msg == null)
             return Scope.NOOP_SCOPE;
 
+        OperationContextAttribute<? extends Message>[] locRegisteredAttrs = registeredAttrs;
+
         assert msg.idBitmap != 0;
-        assert !F.isEmpty(msg.vals);
-        assert msg.vals.length <= MAX_ATTRS_CNT;
+        assert !F.isEmpty(msg.attrs);
+        assert msg.attrs.length <= MAX_ATTRS_CNT;
 
         OperationContext.ContextUpdater updater = OperationContext.ContextUpdater.create();
 
-        for (byte valIdx = 0, maskIdx = 0; valIdx < msg.vals.length; ++valIdx) {
-            Message curVal = msg.vals[valIdx];
+        for (byte valIdx = 0, attrId = 0; valIdx < msg.attrs.length; ++valIdx) {
+            Message curVal = msg.attrs[valIdx];
+
+            while ((msg.idBitmap & (1 << attrId)) == 0)
+                ++attrId;
 
-            while ((msg.idBitmap & (1 << maskIdx)) == 0)
-                ++maskIdx;
+            assert attrId < locRegisteredAttrs.length;
 
-            OperationContextAttribute<Message> attr = (OperationContextAttribute<Message>)attrs.get(maskIdx++);
+            OperationContextAttribute<Message> attr = (OperationContextAttribute<Message>)locRegisteredAttrs[attrId++];
 
             assert attr != null;
 
@@ -139,7 +155,7 @@ public Scope restoreDistributedAttributes(@Nullable OperationContextMessage msg)
     }
 
     /** Restricts further registration of distributed attributes. */
-    public void finishRegistration() {
+    public synchronized void finishRegistration() {
         regFinished = true;
     }
 }

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

@@ -15,10 +15,9 @@
  * limitations under the License.
  */
 
-package org.apache.ignite.internal;
+package org.apache.ignite.internal.thread.context;
 
-import org.apache.ignite.internal.thread.context.OperationContext;
-import org.apache.ignite.internal.thread.context.OperationContextDispatcher;
+import org.apache.ignite.internal.Order;
 import org.apache.ignite.plugin.extensions.communication.Message;
 
 /**
@@ -29,14 +28,20 @@
 public class OperationContextMessage implements Message {
     /** Values of operation context attributes. */
     @Order(0)
-    public Message[] vals;
+    Message[] attrs;
 
     /** Bitmap of effective attributes ids. */
     @Order(1)
-    public byte idBitmap;
+    byte idBitmap;
 
     /** Empty constructor for serialization purposes. */
     public OperationContextMessage() {
         // No-op.
     }
+
+    /** */
+    public OperationContextMessage(byte idBitmap, Message[] attrs) {
+        this.attrs = attrs;
+        this.idBitmap = idBitmap;
+    }
 }

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

@@ -1311,7 +1311,7 @@ private class SocketWriter extends IgniteSpiThread {
          * @param msg Message.
          */
         private void sendMessage(TcpDiscoveryAbstractMessage msg) {
-            msg.opCtxMsg = operationCtxDispatcher.collectDistributedAttributes();
+            msg.opCtxMsg = operationCtxDispatcher.collectDistributedAttributeValues();
 
             synchronized (mux) {
                 queue.add(msg);
@@ -1764,7 +1764,7 @@ private MessageWorker(IgniteLogger log) {
                         ? (TcpDiscoveryAbstractMessage)msg
                         : null;
 
-                    try (Scope ignored = operationCtxDispatcher.restoreDistributedAttributes(dm == null ? null : dm.opCtxMsg)) {
+                    try (Scope ignored = operationCtxDispatcher.restoreRemoteAttributeValues(dm == null ? null : dm.opCtxMsg)) {
                         if (msg instanceof JoinTimeout) {
                             int joinCnt0 = ((JoinTimeout)msg).joinCnt;
 

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

@@ -3021,7 +3021,7 @@ void addMessage(TcpDiscoveryAbstractMessage msg, boolean ignoreHighPriority, boo
             }
 
             if (!fromSocket)
-                msg.opCtxMsg = operationCtxDispatcher.collectDistributedAttributes();
+                msg.opCtxMsg = operationCtxDispatcher.collectDistributedAttributeValues();
 
             if (msg instanceof TraceableMessage tMsg) {
 
@@ -3293,7 +3293,7 @@ else if (msg instanceof TcpDiscoveryAuthFailedMessage)
             if (msg == WAKEUP)
                 return;
 
-            try (Scope ignored = operationCtxDispatcher.restoreDistributedAttributes(msg.opCtxMsg)) {
+            try (Scope ignored = operationCtxDispatcher.restoreRemoteAttributeValues(msg.opCtxMsg)) {
                 processMessage0(msg);
             }
         }
@@ -6186,6 +6186,7 @@ private void processCustomMessage(TcpDiscoveryCustomEventMessage msg, boolean wa
                                 getLocalNodeId(), nextMsg);
 
                             ackMsg.topologyVersion(msg.topologyVersion());
+                            ackMsg.opCtxMsg = operationCtxDispatcher.collectDistributedAttributeValues();
 
                             processCustomMessage(ackMsg, waitForNotification);
                         }

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

@@ -21,8 +21,8 @@
 import java.util.HashSet;
 import java.util.Set;
 import java.util.UUID;
-import org.apache.ignite.internal.OperationContextMessage;
 import org.apache.ignite.internal.Order;
+import org.apache.ignite.internal.thread.context.OperationContextMessage;
 import org.apache.ignite.internal.util.tostring.GridToStringExclude;
 import org.apache.ignite.internal.util.tostring.GridToStringInclude;
 import org.apache.ignite.internal.util.typedef.internal.S;