IGNITE-26608

Author: petrov-mg

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

@@ -66,20 +66,12 @@ T peek() {
     }
 
     /** */
-    ThreadContextSnapshot exportTo(ThreadContextSnapshot snapshot) {
+    ThreadContextSnapshot exportTopTo(ThreadContextSnapshot snapshot) {
         T val = peek();
 
         return val == attr.initialValue() ? snapshot : snapshot.withAttribute(attr, val);
     }
 
-    /** */
-    void restoreInitial(int scopeDepth) {
-        if (isEmpty() || scopedVals.peek().value() == attr.initialValue())
-            return;
-
-        scopedVals.push(new ScopedAttributeValue<>(scopeDepth, attr.initialValue()));
-    }
-
     /** */
     boolean isEmpty() {
         return F.isEmpty(scopedVals);

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

@@ -29,7 +29,7 @@ public class ThreadContextAttributeRegistry {
     private final List<ThreadContextAttribute<?>> attrs = new CopyOnWriteArrayList<>();
 
     /**
-     * Registers attribute with initial value set to {@code null}.
+     * Registers new attribute with initial value set to {@code null}.
      *
      * @see #register(Object)
      */
@@ -38,13 +38,13 @@ public <T> ThreadContextAttribute<T> register() {
     }
 
     /**
-     * Registers attribute with specified initial value. Initial value is returned by
+     * Registers new attribute with specified initial value. Initial value is returned by
      * {@link ThreadContext#get(ThreadContextAttribute)} method if attribute value is not set explicitly.
      * Returned value represents a key used to access attribute value via {@link ThreadContext#get(ThreadContextAttribute)}
      * or {@link ThreadContext#withAttribute(ThreadContextAttribute, Object)} methods.
      *
      * @param initialVal Attribute initial value.
-     * @return Registered attribute instance.
+     * @return Thread Context Attribute instance.
      */
     public synchronized <T> ThreadContextAttribute<T> register(T initialVal) {
         ThreadContextAttribute<T> attr = new ThreadContextAttribute<>(attrs.size(), initialVal);
@@ -56,6 +56,8 @@ public synchronized <T> ThreadContextAttribute<T> register(T initialVal) {
 
     /** */
     <T> ThreadContextAttribute<T> attribute(int id) {
+        assert id < attrs.size();
+
         return (ThreadContextAttribute<T>)attrs.get(id);
     }
 

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

@@ -62,7 +62,7 @@ ThreadContextSnapshot createSnapshot() {
         ThreadContextSnapshot snapshot = ThreadContextSnapshot.emptySnapshot();
 
         for (ScopedAttributeValueStack<?> attrVals : attrs)
-            snapshot = attrVals.exportTo(snapshot);
+            snapshot = attrVals.exportTopTo(snapshot);
 
         return snapshot;
     }
@@ -73,13 +73,18 @@ void restoreSnapshot(ThreadContextSnapshot snapshot) {
             return;
 
         for (int id = attrReg.size() - 1; id >= 0; id--) {
-            if (!snapshot.isEmpty() && snapshot.attribute().id() == id) {
-                put(snapshot.attribute(), snapshot.attributeValue());
+            ThreadContextAttribute<Object> attr = attrReg.attribute(id);
+            Object attrVal;
+
+            if (!snapshot.isEmpty() && snapshot.attributeId() == id) {
+                attrVal = snapshot.attributeValue();
 
                 snapshot = snapshot.previous();
             }
             else
-                attributeValues(id).restoreInitial(activeScopeDepth);
+                attrVal = attr.initialValue();
+
+            put(attr, attrVal);
         }
     }
 
@@ -107,13 +112,13 @@ private void clearActiveScopeData() {
     /** */
     private <T> ScopedAttributeValueStack<T> attributeValues(int id) {
         if (attrs.length <= id)
-            fetchRegisteredAttibutes();
+            fetchRegisteredAttributes();
 
         return (ScopedAttributeValueStack<T>)attrs[id];
     }
 
     /** */
-    private void fetchRegisteredAttibutes() {
+    private void fetchRegisteredAttributes() {
         ScopedAttributeValueStack<?>[] upd = new ScopedAttributeValueStack[attrReg.size()];
 
         if (attrs.length != 0)

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

@@ -20,10 +20,10 @@
 /** */
 public class ThreadContextSnapshot {
     /** */
-    private static final ThreadContextSnapshot EMPTY = new ThreadContextSnapshot(null, null, null);
+    private static final ThreadContextSnapshot EMPTY = new ThreadContextSnapshot(-1, null, null);
 
     /** */
-    private final ThreadContextAttribute<?> attr;
+    private final int attrId;
 
     /** */
     private final Object attrVal;
@@ -32,17 +32,17 @@ public class ThreadContextSnapshot {
     private final ThreadContextSnapshot prev;
 
     /** */
-    private ThreadContextSnapshot(ThreadContextAttribute<?> attr, Object attrVal, ThreadContextSnapshot prev) {
-        this.attr = attr;
+    private ThreadContextSnapshot(int attrId, Object attrVal, ThreadContextSnapshot prev) {
+        this.attrId = attrId;
         this.attrVal = attrVal;
         this.prev = prev;
     }
 
     /** */
-    <T> ThreadContextAttribute<T> attribute() {
+    int attributeId() {
         assert !isEmpty();
 
-        return (ThreadContextAttribute<T>)attr;
+        return attrId;
     }
 
     /** */
@@ -66,7 +66,7 @@ boolean isEmpty() {
 
     /** */
     <T> ThreadContextSnapshot withAttribute(ThreadContextAttribute<T> attr, T val) {
-        return new ThreadContextSnapshot(attr, val, this);
+        return new ThreadContextSnapshot(attr.id(), val, this);
     }
 
     /** */