IGNITE-28915 Fixed Operation Context propagation for postponed messages

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

@@ -113,7 +113,7 @@ public static <T1, T2> Scope set(
         OperationContextAttribute<T1> attr1, T1 val1,
         OperationContextAttribute<T2> attr2, T2 val2
     ) {
-        return ContextUpdater.create().set(attr1, val1).set(attr2, val2).apply();
+        return Updater.create().set(attr1, val1).set(attr2, val2).apply();
     }
 
     /**
@@ -136,7 +136,7 @@ public static <T1, T2, T3> Scope set(
         OperationContextAttribute<T2> attr2, T2 val2,
         OperationContextAttribute<T3> attr3, T3 val3
     ) {
-        return ContextUpdater.create().set(attr1, val1).set(attr2, val2).set(attr3, val3).apply();
+        return Updater.create().set(attr1, val1).set(attr2, val2).set(attr3, val3).apply();
     }
 
     /**
@@ -322,49 +322,88 @@ private static class AttributeValueHolder<T> {
     }
 
     /** Allows to change multiple attribute values in a single update operation and skip updates that changes nothing. */
-    static class ContextUpdater {
+    static class Updater extends AttributeCollector {
         /** */
-        private static final int INIT_UPDATES_CAPACITY = 3;
+        private Updater(OperationContext ctx) {
+            super(ctx);
+        }
 
-        /** */
-        private final OperationContext ctx;
+        /** */  
+        <T> Updater set(OperationContextAttribute<T> attr, T val) {
+            if (ctx.getInternal(attr) == val)
+                return this;
+
+            add(attr, val);
+
+            return this;
+        }
 
         /** */
-        private List<AttributeValueHolder<?>> updates;
+        Scope apply() {
+            return isEmpty() ? NOOP_SCOPE : ctx.applyAttributeUpdates(toArray());
+        }
 
         /** */
-        private ContextUpdater(OperationContext ctx) {
-            this.ctx = ctx;
+        static Updater create() {
+            return new Updater(INSTANCE.get());
         }
+    }
 
-        /** */  
-        <T> ContextUpdater set(OperationContextAttribute<T> attr, T val) {
-            if (ctx.getInternal(attr) == val)
-                return this;
+    /** */
+    static class Restorer extends AttributeCollector {
+        /** */
+        private Restorer(OperationContext ctx) {
+            super(ctx);
+        }
 
-            if (updates == null)
-                updates = new ArrayList<>(INIT_UPDATES_CAPACITY);
+        /** */
+        Scope restore() {
+            return ctx.restoreSnapshotInternal(isEmpty() ? null : ctx.new Update(toArray(), null));
+        }
 
-            updates.add(new AttributeValueHolder<>(attr, val));
+        /** */
+        static Restorer create() {
+            return new Restorer(INSTANCE.get());
+        }
 
-            return this;
+        /** */
+        static Scope restoreEmpty() {
+            return INSTANCE.get().restoreSnapshotInternal(null);
         }
+    }
 
+    /** */
+    private abstract static class AttributeCollector {
         /** */
-        Scope apply() {
-            if (F.isEmpty(updates))
-                return NOOP_SCOPE;
+        private static final int INIT_CAPACITY = 3;
 
-            AttributeValueHolder<?>[] sealedUpdates = new AttributeValueHolder[updates.size()];
+        /** */
+        protected final OperationContext ctx;
 
-            updates.toArray(sealedUpdates);
+        /** */
+        private List<AttributeValueHolder<?>> attrs;
 
-            return ctx.applyAttributeUpdates(sealedUpdates);
+        /** */
+        protected AttributeCollector(OperationContext ctx) {
+            this.ctx = ctx;
+        }
+
+        /** */
+        <T> void add(OperationContextAttribute<T> attr, T val) {
+            if (attrs == null)
+                attrs = new ArrayList<>(INIT_CAPACITY);
+
+            attrs.add(new AttributeValueHolder<>(attr, val));
+        }
+
+        /** */
+        boolean isEmpty() {
+            return F.isEmpty(attrs);
         }
 
         /** */
-        static ContextUpdater create() {
-            return new ContextUpdater(INSTANCE.get());
+        protected AttributeValueHolder<?>[] toArray() {
+            return attrs.toArray(AttributeValueHolder<?>[]::new);
         }
     }
 }

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

@@ -259,7 +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.thread.context.OperationContextSnapshotMessage;
 import org.apache.ignite.internal.util.GridByteArrayList;
 import org.apache.ignite.internal.util.GridIntList;
 import org.apache.ignite.internal.util.GridPartitionStateMap;
@@ -712,7 +712,7 @@ public CoreMessagesProvider(Marshaller dfltMarsh, Marshaller schemaAwareMarsh, C
 
         // [13400 - 13500]: Operation context messages.
         msgIdx = 13400;
-        withNoSchema(OperationContextMessage.class);
+        withNoSchema(OperationContextSnapshotMessage.class);
         withNoSchema(SecurityContextWrapper.class);
 
         // [13600 - 13700]: Rolling Upgrade messages.

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

@@ -451,7 +451,7 @@ public void resetMetrics() {
                 try {
                     GridIoMessage msg0 = (GridIoMessage)msg;
 
-                    try (Scope ignored = ctx.operationContextDispatcher().restoreRemoteAttributeValues(msg0.opCtxMsg)) {
+                    try (Scope ignored = ctx.operationContextDispatcher().restoreSnapshot(msg0.opCtxSnp)) {
                         onMessage0(nodeId, msg0, msgC);
                     }
                 }
@@ -2027,13 +2027,7 @@ public GridIoMessage createGridIoMessage(
         long timeout,
         boolean skipOnTimeout
     ) {
-        GridIoMessage res;
-
-        res = new GridIoMessage(plc, topic, msg, ordered, timeout, skipOnTimeout);
-
-        res.opCtxMsg = ctx.operationContextDispatcher().collectDistributedAttributeValues();
-
-        return res;
+        return new GridIoMessage(plc, topic, msg, ordered, timeout, skipOnTimeout, ctx.operationContextDispatcher().createSnapshot());
     }
 
     /**
@@ -3782,12 +3776,14 @@ void unwind(GridMessageListener lsnr) {
             assert reserved.get();
 
             for (OrderedMessageContainer mc = msgs.poll(); mc != null; mc = msgs.poll()) {
-                try {
-                    invokeListener(plc, lsnr, nodeId, mc.message.message());
-                }
-                finally {
-                    if (mc.closure != null)
-                        mc.closure.run();
+                try (Scope ignored0 = ctx.operationContextDispatcher().restoreSnapshot(mc.message.opCtxSnp)) {
+                    try {
+                        invokeListener(plc, lsnr, nodeId, mc.message.message());
+                    }
+                    finally {
+                        if (mc.closure != null)
+                            mc.closure.run();
+                    }
                 }
             }
         }

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

@@ -22,7 +22,7 @@
 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.thread.context.OperationContextMessage;
+import org.apache.ignite.internal.thread.context.OperationContextSnapshotMessage;
 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;
@@ -64,7 +64,7 @@ public class GridIoMessage implements Message, MessageWrapper {
     /** Effective operation context attributes to propagate. */
     @Order(6)
     @GridToStringInclude
-    public @Nullable OperationContextMessage opCtxMsg;
+    @Nullable OperationContextSnapshotMessage opCtxSnp;
 
     /**
      * Default constructor.
@@ -80,14 +80,16 @@ public GridIoMessage() {
      * @param ordered Message ordered flag.
      * @param timeout Timeout.
      * @param skipOnTimeout Whether message can be skipped on timeout.
+     * @param opCtxSnp Operation Context snapshot.
      */
     public GridIoMessage(
         byte plc,
         Object topic,
         Message msg,
         boolean ordered,
         long timeout,
-        boolean skipOnTimeout
+        boolean skipOnTimeout,
+        @Nullable OperationContextSnapshotMessage opCtxSnp
     ) {
         assert topic != null;
         assert msg != null;
@@ -98,6 +100,7 @@ public GridIoMessage(
         this.ordered = ordered;
         this.timeout = timeout;
         this.skipOnTimeout = skipOnTimeout;
+        this.opCtxSnp = opCtxSnp;
     }
 
     /**

modules/core/src/main/java/org/apache/ignite/internal/processors/pool/PoolProcessor.java

@@ -273,7 +273,7 @@ public PoolProcessor(GridKernalContext ctx) {
                         throw new IgniteException("Failed to register IO executor pool because its ID as " +
                             "already used: " + id);
 
-                    extPools[id] = ctx.security().enabled() ? OperationContextAwareIoPool.wrap(ex) : ex;
+                    extPools[id] = OperationContextAwareIoPool.wrap(ex);
                 }
             }
         }

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.DistributedAttributeRegistry.SECURITY;
+import static org.apache.ignite.internal.thread.context.DistributedAttributeIdRegistry.SECURITY;
 import static org.apache.ignite.plugin.security.SecurityPermission.ADMIN_USER_ACCESS;
 import static org.apache.ignite.plugin.security.SecurityPermission.JOIN_AS_SERVER;
 

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.DistributedAttributeRegistry;
+import org.apache.ignite.internal.thread.context.DistributedAttributeIdRegistry;
 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#collectDistributedAttributeValues()
- * @see DistributedAttributeRegistry#SECURITY
+ * @see OperationContextDispatcher#createSnapshot()
+ * @see DistributedAttributeIdRegistry#SECURITY
  */
 public class SecurityContextWrapper implements Message {
     /** A value of {@link SecuritySubject#id()} */

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

@@ -23,7 +23,7 @@
  * 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 class DistributedAttributeIdRegistry {
+    /** ID Reserved for {@link SecurityContext} propagation. */
     public static final byte SECURITY = 0;
 }

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

@@ -16,9 +16,7 @@
  */
 package org.apache.ignite.internal.thread.context;
 
-import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.List;
 import org.apache.ignite.IgniteException;
 import org.apache.ignite.internal.util.typedef.F;
 import org.apache.ignite.plugin.extensions.communication.Message;
@@ -41,7 +39,7 @@
  * {@link #MAX_ATTRS_CNT} for implementation reasons.</p>
  *
  * @see OperationContext
- * @see OperationContextMessage
+ * @see OperationContextSnapshotMessage
  */
 public class OperationContextDispatcher {
     /** Maximal number of supported distributed attributes. */
@@ -62,7 +60,7 @@ public class OperationContextDispatcher {
      * <p>Registered attribute value is automatically captured and propagated between cluster nodes
      * during the messages transmission.</p>
      *
-     * @see DistributedAttributeRegistry
+     * @see DistributedAttributeIdRegistry
      */
     public synchronized <T extends Message> void registerDistributedAttribute(int id, OperationContextAttribute<T> attr) {
         if (regFinished)
@@ -89,14 +87,13 @@ public synchronized <T extends Message> void registerDistributedAttribute(int id
      *
      * @see OperationContext#get(OperationContextAttribute)
      */
-    public @Nullable OperationContextMessage collectDistributedAttributeValues() {
+    public @Nullable OperationContextSnapshotMessage createSnapshot() {
         OperationContextAttribute<? extends Message>[] locRegisteredAttrs = registeredAttrs;
 
         if (locRegisteredAttrs.length == 0)
             return null;
 
-        byte bitmap = 0;
-        List<Message> vals = null;
+        OperationContextSnapshotMessage.Builder snpBuilder = OperationContextSnapshotMessage.Builder.create();
 
         for (int id = 0; id < locRegisteredAttrs.length; id++) {
             OperationContextAttribute<? extends Message> attr = locRegisteredAttrs[id];
@@ -106,40 +103,30 @@ public synchronized <T extends Message> void registerDistributedAttribute(int id
 
             Message curVal = OperationContext.get(attr);
 
-            if (curVal == attr.initialValue())
-                continue;
-
-            if (vals == null)
-                vals = new ArrayList<>(MAX_ATTRS_CNT / 2);
-
-            byte mask = (byte)(1 << id);
-
-            assert (bitmap & mask) == 0;
-
-            vals.add(curVal);
-            bitmap |= mask;
+            if (curVal != attr.initialValue())
+                snpBuilder.add(id, curVal);
         }
 
-        return bitmap == 0 ? null : new OperationContextMessage(bitmap, vals.toArray(Message[]::new));
+        return snpBuilder.isEmpty() ? null : snpBuilder.build();
     }
 
-    /** Restores distributed {@link OperationContextAttribute} values received from a remote node. */
-    public Scope restoreRemoteAttributeValues(@Nullable OperationContextMessage msg) {
-        if (msg == null)
-            return Scope.NOOP_SCOPE;
+    /** Restores {@link OperationContextAttribute} values received from a remote node. */
+    public Scope restoreSnapshot(@Nullable OperationContextSnapshotMessage snp) {
+        if (snp == null)
+            return OperationContext.Restorer.restoreEmpty();
 
         OperationContextAttribute<? extends Message>[] locRegisteredAttrs = registeredAttrs;
 
-        assert msg.idBitmap != 0;
-        assert !F.isEmpty(msg.attrs);
-        assert msg.attrs.length <= MAX_ATTRS_CNT;
+        assert snp.idBitmap != 0;
+        assert !F.isEmpty(snp.attrs);
+        assert snp.attrs.length <= MAX_ATTRS_CNT;
 
-        OperationContext.ContextUpdater updater = OperationContext.ContextUpdater.create();
+        OperationContext.Restorer ctxRestorer = OperationContext.Restorer.create();
 
-        for (byte valIdx = 0, attrId = 0; valIdx < msg.attrs.length; ++valIdx) {
-            Message curVal = msg.attrs[valIdx];
+        for (byte valIdx = 0, attrId = 0; valIdx < snp.attrs.length; ++valIdx) {
+            Message attrVal = snp.attrs[valIdx];
 
-            while ((msg.idBitmap & (1 << attrId)) == 0)
+            while ((snp.idBitmap & (1 << attrId)) == 0)
                 ++attrId;
 
             assert attrId < locRegisteredAttrs.length;
@@ -148,10 +135,10 @@ public Scope restoreRemoteAttributeValues(@Nullable OperationContextMessage msg)
 
             assert attr != null;
 
-            updater.set(attr, curVal);
+            ctxRestorer.add(attr, attrVal);
         }
 
-        return updater.apply();
+        return ctxRestorer.restore();
     }
 
     /** Restricts further registration of distributed attributes. */