IGNITE-26608

Author: petrov-mg

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

@@ -0,0 +1,98 @@
+/*
+ * 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;
+
+/** Represents mapping of {@link ContextAttribute} to their corresponding values with the ability to be attached to the thread. */
+public final class Context {
+    /**
+     * Creates a new {@link AttributeValueHolder} with a single mapping of the specified {@link AttributeValueHolder}
+     * to the specified value. The {@link AttributeValueHolder} can be used to accumulate mappings.
+     *
+     * @see AttributeValueHolder#with(ContextAttribute, Object)
+     */
+    public static <T> AttributeValueHolder with(ContextAttribute<T> attr, T val) {
+        return AttributeValueHolder.ROOT.with(attr, val);
+    }
+
+    /** */
+    public static final class AttributeValueHolder extends ContextDataChain<AttributeValueHolder> {
+        /** */
+        private static final AttributeValueHolder ROOT = new AttributeValueHolder();
+
+        /** */
+        private final ContextAttribute<?> attr;
+
+        /** */
+        private final Object val;
+
+        /** */
+        private AttributeValueHolder() {
+            attr = null;
+            val = null;
+        }
+
+        /** */
+        private <T> AttributeValueHolder(ContextAttribute<T> attr, T val, AttributeValueHolder prev) {
+            super(prev.storedAttributeIdBits() | attr.bitmask(), prev);
+
+            this.attr = attr;
+            this.val = val;
+        }
+
+        /** Adds a new mapping of the specified attribute to its value to the current {@link ContextDataChain}. */
+        public <T> AttributeValueHolder with(ContextAttribute<T> attr, T val) {
+            return attr.get() == val ? this : new AttributeValueHolder(attr, val, this);
+        }
+
+        /** */
+        ContextAttribute<?> attribute() {
+            assert !isEmpty();
+
+            return attr;
+        }
+
+        /** */
+        <T> T value() {
+            assert !isEmpty();
+
+            return (T)val;
+        }
+
+        /** {@inheritDoc} */
+        @Override boolean isEmpty() {
+            return this == ROOT;
+        }
+
+        /**
+         * Attaches {@link ContextAttribute} values stored in current {@link ContextDataChain} to the thread
+         * this method is called from. If {@link ContextAttribute} value was already attached for the current thread,
+         * its value will be stashed and replaced by the new ones.
+         *
+         * @return {@link Scope} instance that, when closed, resets the values for all {@link ContextAttribute}s added
+         * to the current Context and restores them to the previously attached values, if any.
+         */
+        public Scope attach() {
+            if (isEmpty())
+                return Scope.NOOP_SCOPE;
+
+            ThreadLocalContextStorage.get().attach(this);
+
+            return () -> ThreadLocalContextStorage.get().detach(this);
+        }
+    }
+}

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

@@ -19,13 +19,17 @@
 
 import java.util.concurrent.atomic.AtomicInteger;
 
-/** */
+/**
+ * Represents an accessor to Context Attribute value bound to the thread.
+ *
+ * @see Context
+ */
 public class ContextAttribute<T> {
     /** */
     static final AtomicInteger ID_GEN = new AtomicInteger();
 
     /** */
-    static final int MAX_ATTRIBUTE_CNT = Integer.SIZE;
+    static final int MAX_ATTR_CNT = Integer.SIZE;
 
     /** */
     private final byte id;
@@ -53,32 +57,41 @@ int bitmask() {
         return bitmask;
     }
 
-    /** */
+    /**
+     * Gets the value of the Context Attribute bound to the thread from which this method is called. 
+     *  
+     * @see Context#with(ContextAttribute, Object) 
+     */
     public T get() {
-        ScopedContext sc = ThreadLocalContextStorage.get().findScopedContextFor(this);
+        Context.AttributeValueHolder valHolder = ThreadLocalContextStorage.get().findValueHolderFor(this);
 
-        return sc == null ? initVal : sc.value();
+        return valHolder == null ? initVal : valHolder.value();
     }
 
     /** */
+    static int highReservedId() {
+        return ID_GEN.get();
+    }
+
+    /** Creates new instance of the Context Attribute with Initial Value set to {@code null}. */
     public static <T> ContextAttribute<T> newInstance() {
         return newInstance(null);
     }
 
-    /** */
+    /**
+     * Creates new instance of the Context Attribute with the specified Initial Value. The Initial Value is returned
+     * by {@link ContextAttribute#get} method if the Attribute's value is not explicitly set in the Context.
+     *
+     * <p>
+     * Note, that the maximum number of attribute instances that can be created is currently limited to
+     * {@link #MAX_ATTR_CNT} for implementation reasons.
+     * </p>
+     */
     public static <T> ContextAttribute<T> newInstance(T initVal) {
         int id = ID_GEN.getAndIncrement();
 
-        if (MAX_ATTRIBUTE_CNT <= id) {
-            throw new RuntimeException("Exceeded maximum supported number of created Context Attributes instances" +
-                " [maxCnt=" + MAX_ATTRIBUTE_CNT + ']');
-        }
+        assert id < MAX_ATTR_CNT : "Exceeded maximum supported number of created Attributes instances [maxCnt=" + MAX_ATTR_CNT + ']';
 
         return new ContextAttribute<>((byte)id, initVal);
     }
-
-    /** */
-    static int highReservedId() {
-        return ID_GEN.get();
-    }
 }

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

@@ -20,34 +20,34 @@
 /** */
 public abstract class ContextDataChain<T> {
     /** */
-    private final int storedAttrBits;
+    private final int storedAttrIdBits;
 
     /** */
     private final T prev;
 
     /** */
     protected ContextDataChain() {
-        storedAttrBits = 0;
+        storedAttrIdBits = 0;
         prev = null;
     }
 
     /** */
-    protected ContextDataChain(int storedAttrBits, T prev) {
-        this.storedAttrBits = storedAttrBits;
+    protected ContextDataChain(int storedAttrIdBits, T prev) {
+        this.storedAttrIdBits = storedAttrIdBits;
         this.prev = prev;
     }
 
     /** */
     abstract boolean isEmpty();
 
     /** */
-    int storedAttributeBits() {
-        return storedAttrBits;
+    int storedAttributeIdBits() {
+        return storedAttrIdBits;
     }
 
     /** */
     boolean containsValueFor(ContextAttribute<?> attr) {
-        return (storedAttrBits & attr.bitmask()) != 0;
+        return (storedAttrIdBits & attr.bitmask()) != 0;
     }
 
     /** */

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

@@ -23,18 +23,18 @@ public class ContextSnapshot extends ContextDataChain<ContextSnapshot> {
     static final ContextSnapshot ROOT = new ContextSnapshot();
 
     /** */
-    private final ScopedContext scopedCtx;
+    private final Context.AttributeValueHolder data;
 
     /** */
     private ContextSnapshot() {
-        scopedCtx = null;
+        data = null;
     }
 
     /** */
-    private ContextSnapshot(ScopedContext scopedCtx, ContextSnapshot prev) {
-        super(scopedCtx.storedAttributeBits() | prev.storedAttributeBits(), prev);
+    private ContextSnapshot(Context.AttributeValueHolder data, ContextSnapshot prev) {
+        super(data.storedAttributeIdBits() | prev.storedAttributeIdBits(), prev);
 
-        this.scopedCtx = scopedCtx;
+        this.data = data;
     }
 
     /** */
@@ -55,15 +55,15 @@ public Scope restore() {
     }
 
     /** */
-    ContextSnapshot attach(ScopedContext scopedCtx) {
-        return new ContextSnapshot(scopedCtx, this);
+    ContextSnapshot attach(Context.AttributeValueHolder data) {
+        return new ContextSnapshot(data, this);
     }
 
     /** */
-    ScopedContext scopedContext() {
+    Context.AttributeValueHolder data() {
         assert !isEmpty();
 
-        return scopedCtx;
+        return data;
     }
 
     /** {@inheritDoc} */

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

@@ -17,11 +17,14 @@
 
 package org.apache.ignite.internal.thread.context;
 
-/** */
+/**
+ * Represents an arbitrary Scope. A Scope is active from the moment it is created until the {@link #close()} method is
+ * called on the {@link Scope} instance. It is strongly encouraged to use a try-with-resources block to close a Scope.
+ */
 public interface Scope extends AutoCloseable {
-    /** */
+    /** Scope instance that does nothing when closed. */
     Scope NOOP_SCOPE = () -> {};
 
-    /** {@inheritDoc} */
+    /** Closes the scope. This operation cannot fail. */
     @Override void close();
 }