IGNITE-28770 : Introduce a operation context attributes registry v2

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 - 13600]: 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,42 @@
+/*
+ * 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 org.apache.ignite.internal.thread.context.DistributedOperationContextManager;
+import org.apache.ignite.internal.thread.context.OperationContext;
+import org.apache.ignite.plugin.extensions.communication.Message;
+
+/**
+ * Transport for {@link OperationContext} distributed attributes.
+ *
+ * @see DistributedOperationContextManager
+ */
+public class DistributedOperationContextMessage implements Message {
+    /** Values of operation context attributes. */
+    @Order(0)
+    public Message[] vals;
+
+    /** Bitmap 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/thread/context/DistributedOperationContextManager.java

@@ -0,0 +1,128 @@
+/*
+ * 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.List;
+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;
+
+/**
+ * A manager of {@link OperationContextAttribute} which are required to be propagated through a cluster.
+ * Has own attributes ids compared to a local node's {@link OperationContext}.
+ *
+ * @see OperationContext
+ * @see DistributedOperationContextMessage
+ */
+public class DistributedOperationContextManager {
+    /** */
+    private static final DistributedOperationContextManager INSTANCE = new DistributedOperationContextManager();
+
+    /** Maximal number of supported distributed attributes. */
+    static final byte MAX_DISTRIBUTED_ATTR_CNT = Byte.SIZE;
+
+    /** Registered distributed attributes by their cluster-wide id. */
+    private final Map<Byte, OperationContextAttribute<? extends Message>> attrs = new ConcurrentSkipListMap<>();
+
+    /** */
+    public static DistributedOperationContextManager instance() {
+        return INSTANCE;
+    }
+
+    /**
+     * Creates and registers a distributable {@link OperationContextAttribute}.
+     *
+     * @param id Cluster-wide id of a distributed operation context attribute.
+     * @param initVal The attribute's unitial value.
+     */
+    public <T extends Message> OperationContextAttribute<T> createDistributedAttribute(byte id, @Nullable T initVal) {
+        assert id >= 0 && id < MAX_DISTRIBUTED_ATTR_CNT : "Invalid distributed attributed id [id=" + id + ']';
+
+        return (OperationContextAttribute<T>)attrs.compute(id, (id0, attr0) -> {
+            if (attr0 != null)
+                throw new IgniteException("Duplicated distributed attribute id [id=" + id + ']');
+
+            return OperationContextAttribute.newInstance(initVal);
+        });
+    }
+
+    /**
+     * Requests current {@link OperationContext} for its effective attributes and collects ones which are also registered
+     * as distbibued attributes.
+     *
+     * @return A message to send current effective distributed attributes. {@code null}, if there are no
+     * effective attributes in {@link OperationContext} or none of them is a distributed attribute.
+     */
+    public @Nullable DistributedOperationContextMessage collectDistributedAttributes() {
+        DistributedOperationContextMessage res = null;
+        List<Message> vals = null;
+
+        for (Map.Entry<Byte, OperationContextAttribute<? extends Message>> e : attrs.entrySet()) {
+            OperationContextAttribute<? extends Message> attr = e.getValue();
+
+            Message curVal = OperationContext.get(attr);
+
+            if (curVal != attr.initialValue()) {
+                if (res == null) {
+                    res = new DistributedOperationContextMessage();
+
+                    vals = new ArrayList<>(MAX_DISTRIBUTED_ATTR_CNT / 2);
+                }
+
+                byte mask = (byte)(1 << e.getKey());
+
+                assert (res.idBitmap & mask) == 0;
+
+                vals.add(curVal);
+                res.idBitmap |= mask;
+            }
+        }
+
+        if (res != null)
+            res.vals = vals.toArray(vals.toArray(new Message[vals.size()]));
+
+        return res;
+    }
+
+    /** Sets the received distributed operation context attributes (if any) into current {@link OperationContext}. */
+    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.length <= MAX_DISTRIBUTED_ATTR_CNT;
+
+        OperationContext.ContextUpdater updater = OperationContext.ContextUpdater.create();
+
+        for (byte valIdx = 0, maskIdx = 0; valIdx < msg.vals.length; ++valIdx) {
+            Message curVal = msg.vals[valIdx];
+
+            while ((msg.idBitmap & (1 << maskIdx)) == 0)
+                ++maskIdx;
+
+            updater.set((OperationContextAttribute<Message>)attrs.get(maskIdx++), curVal);
+        }
+
+        return updater.apply();
+    }
+}