Skip to content

Commit 5c25680

Browse files
committed
1 parent 7f53ad0 commit 5c25680

File tree

6 files changed

+143
-85
lines changed

6 files changed

+143
-85
lines changed

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,7 @@
1818
package org.apache.ignite.internal.thread.context;
1919

2020
/** */
21-
22-
/** */
23-
public class AttributeValueHolder {
21+
public final class AttributeValueHolder {
2422
/** */
2523
private final ContextAttribute<?> attr;
2624

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,11 @@ int storedAttributeIdBits() {
9696
return storedAttrIdBits;
9797
}
9898

99+
/** */
100+
public static ContextSnapshot createSnapshot() {
101+
return ThreadLocalContextStorage.get().createSnapshot();
102+
}
103+
99104
/** */
100105
public static final class Builder {
101106
/** */

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ protected static <T> T wrap(T delegate, BiFunction<T, ContextSnapshot, T> wrappe
4949
if (delegate == null || delegate instanceof ContextAwareWrapper)
5050
return delegate;
5151

52-
ContextSnapshot snapshot = ContextSnapshot.capture();
52+
ContextSnapshot snapshot = Context.createSnapshot();
5353

5454
if (ignoreEmptyContext && snapshot.isEmpty())
5555
return delegate;

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

Lines changed: 14 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -17,26 +17,21 @@
1717

1818
package org.apache.ignite.internal.thread.context;
1919

20-
import java.util.ArrayList;
21-
import java.util.Collection;
22-
import java.util.Collections;
23-
import org.apache.ignite.internal.util.typedef.F;
24-
2520
/**
26-
* Represents Snapshot of all {@link ContextAttribute}s and their corresponding values. Note that Snapshot also stores
27-
* the states of {@link ContextAttribute}s for which value are note explicitly specified.
21+
* Represents Snapshot of all {@link ContextAttribute}s and their corresponding values. Note that Snapshot also must
22+
* store he states of {@link ContextAttribute}s for which value are note explicitly specified.
2823
*/
29-
public class ContextSnapshot {
30-
/** */
31-
private static final ContextSnapshot EMPTY = new ContextSnapshot(Collections.emptyList());
32-
24+
public interface ContextSnapshot {
3325
/** */
34-
private final Collection<Context> ctxStackSnp;
26+
ContextSnapshot EMPTY = new ContextSnapshot() {
27+
@Override public Scope restore() {
28+
return Scope.NOOP_SCOPE;
29+
}
3530

36-
/** */
37-
private ContextSnapshot(Collection<Context> ctxStackSnp) {
38-
this.ctxStackSnp = new ArrayList<>(ctxStackSnp);
39-
}
31+
@Override public boolean isEmpty() {
32+
return true;
33+
}
34+
};
4035

4136
/**
4237
* Stashes all {@link ContextAttribute} values attached to the thread from which this method is called and replaces
@@ -45,28 +40,8 @@ private ContextSnapshot(Collection<Context> ctxStackSnp) {
4540
* @return {@link Scope} instance that, when closed, restores the values of all {@link ContextAttribute}s to the
4641
* state they were in before the current method was called.
4742
*/
48-
public Scope restore() {
49-
ThreadLocalContextStorage threadStorage = ThreadLocalContextStorage.get();
50-
51-
Collection<Context> prev = threadStorage.contextStackSnapshot();
52-
53-
threadStorage.reinitialize(ctxStackSnp);
54-
55-
return () -> ThreadLocalContextStorage.get().reinitialize(prev);
56-
}
57-
58-
/** */
59-
public boolean isEmpty() {
60-
return F.isEmpty(ctxStackSnp);
61-
}
62-
63-
/**
64-
* Captures Snapshot of all {@link ContextAttribute}s and their corresponding values attached to the thread from which
65-
* this method is called.
66-
*/
67-
public static ContextSnapshot capture() {
68-
Collection<Context> ctxStackSnp = ThreadLocalContextStorage.get().contextStackSnapshot();
43+
public Scope restore();
6944

70-
return ctxStackSnp.isEmpty() ? EMPTY : new ContextSnapshot(ctxStackSnp);
71-
}
45+
/** @return Whether snapshot is empty. */
46+
public boolean isEmpty();
7247
}

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

Lines changed: 73 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,7 @@
1717

1818
package org.apache.ignite.internal.thread.context;
1919

20-
import java.util.ArrayList;
2120
import java.util.Arrays;
22-
import java.util.Collection;
23-
import java.util.Collections;
24-
import java.util.LinkedList;
25-
import java.util.List;
2621
import org.jetbrains.annotations.Nullable;
2722

2823
/** Represents storage of {@link ContextAttribute}s and their values bound to the thread. */
@@ -31,10 +26,7 @@ class ThreadLocalContextStorage {
3126
private static final ThreadLocal<ThreadLocalContextStorage> INSTANCE = ThreadLocal.withInitial(ThreadLocalContextStorage::new);
3227

3328
/** */
34-
private final LinkedList<Context> ctxStack = new LinkedList<>();
35-
36-
/** */
37-
private final LinkedList<Integer> storedAttrIdBitsStack = new LinkedList<>();
29+
private ContextStackNode ctxStackHead;
3830

3931
/** */
4032
private final Cache cache = new Cache();
@@ -62,64 +54,55 @@ private ThreadLocalContextStorage() {
6254

6355
/** */
6456
void attach(Context ctx) {
65-
ctxStack.push(ctx);
66-
67-
storedAttrIdBitsStack.push(storedAttrIdBitsStack.isEmpty()
68-
? ctx.storedAttributeIdBits()
69-
: ctx.storedAttributeIdBits() | storedAttrIdBitsStack.peek()
57+
ctxStackHead = new ContextStackNode(
58+
ctx,
59+
ctxStackHead == null ? ctx.storedAttributeIdBits() : ctxStackHead.storedAttrIdBits | ctx.storedAttributeIdBits(),
60+
ctxStackHead
7061
);
7162

7263
storeValuesToCache(ctx);
7364
}
7465

7566
/** */
7667
void detach(Context ctx) {
77-
Context top = ctxStack.pop();
68+
assert ctxStackHead != null : "Scopes must be closed in the same order and in the same thread they are opened";
69+
70+
Context top = ctxStackHead.ctx;
7871

7972
assert top == ctx : "Scopes must be closed in the same order and in the same thread they are opened";
8073

81-
storedAttrIdBitsStack.pop();
74+
ctxStackHead = ctxStackHead.prev;
8275

8376
cache.invalidateByIndexBits(top.storedAttributeIdBits());
8477
}
8578

8679
/** */
87-
Collection<Context> contextStackSnapshot() {
88-
if (ctxStack.isEmpty())
89-
return Collections.emptyList();
90-
91-
List<Context> res = new ArrayList<>(ctxStack.size());
92-
93-
ctxStack.descendingIterator().forEachRemaining(res::add);
94-
95-
return res;
80+
ContextSnapshot createSnapshot() {
81+
return ctxStackHead == null ? ContextSnapshot.EMPTY : new ThreadLocalContextSnapshot(ctxStackHead);
9682
}
9783

9884
/** */
99-
void reinitialize(Collection<Context> ctxStackSnp) {
100-
ctxStack.clear();
101-
storedAttrIdBitsStack.clear();
85+
private void restoreStackState(ContextStackNode ctxStackHead) {
86+
this.ctxStackHead = ctxStackHead;
10287

10388
cache.invalidate();
104-
105-
ctxStackSnp.forEach(this::attach);
10689
}
10790

10891
/** */
10992
private boolean containsValueFor(ContextAttribute<?> attr) {
110-
if (storedAttrIdBitsStack.isEmpty())
93+
if (ctxStackHead == null)
11194
return false;
11295

113-
return (storedAttrIdBitsStack.peek() & attr.bitmask()) != 0;
96+
return (ctxStackHead.storedAttrIdBits & attr.bitmask()) != 0;
11497
}
11598

11699
/** */
117100
private AttributeValueHolder searchFor(ContextAttribute<?> attr) {
118-
for (Context ctx : ctxStack) {
119-
if (!ctx.containsValueFor(attr))
101+
for (ContextStackNode stackNode = ctxStackHead; stackNode != null; stackNode = stackNode.prev) {
102+
if (!stackNode.ctx.containsValueFor(attr))
120103
continue;
121104

122-
for (AttributeValueHolder valHolder : ctx) {
105+
for (AttributeValueHolder valHolder : stackNode.ctx) {
123106
if (valHolder.attribute().id() == attr.id()) {
124107
cache.store(valHolder);
125108

@@ -140,7 +123,7 @@ private void storeValuesToCache(Context ctx) {
140123
for (AttributeValueHolder h : ctx) {
141124
ContextAttribute<?> attr = h.attribute();
142125

143-
// The value for an attribute can be specified multiple times during context creation - we ignore all but the last one.
126+
// The value for an attribute can be specified multiple times during context creation - we ignore all but the first one.
144127
if ((processed & attr.bitmask()) != 0)
145128
continue;
146129

@@ -214,4 +197,58 @@ private void grow(int size) {
214197
vals = upd;
215198
}
216199
}
200+
201+
/** */
202+
private static class ContextStackNode {
203+
/** */
204+
final Context ctx;
205+
206+
/** */
207+
final int storedAttrIdBits;
208+
209+
/** */
210+
final ContextStackNode prev;
211+
212+
/** */
213+
ContextStackNode(Context ctx, int storedAttrIdBits, ContextStackNode prev) {
214+
this.ctx = ctx;
215+
this.storedAttrIdBits = storedAttrIdBits;
216+
this.prev = prev;
217+
}
218+
}
219+
220+
/** */
221+
private static class ThreadLocalContextSnapshot implements ContextSnapshot {
222+
/** */
223+
private final ContextStackNode ctxStackHead;
224+
225+
/** */
226+
private ThreadLocalContextSnapshot(ContextStackNode ctxStackHead) {
227+
assert ctxStackHead != null;
228+
229+
this.ctxStackHead = ctxStackHead;
230+
}
231+
232+
/** {@inheritDoc} */
233+
@Override public Scope restore() {
234+
ThreadLocalContextStorage threadStorage = get();
235+
236+
ContextStackNode stash = threadStorage.ctxStackHead;
237+
238+
threadStorage.restoreStackState(ctxStackHead);
239+
240+
return () -> {
241+
ThreadLocalContextStorage ts = get();
242+
243+
assert ts.ctxStackHead == ctxStackHead : "Scopes must be closed in the same order and in the same thread they are opened";
244+
245+
ts.restoreStackState(stash);
246+
};
247+
}
248+
249+
/** {@inheritDoc} */
250+
@Override public boolean isEmpty() {
251+
return false;
252+
}
253+
}
217254
}

modules/core/src/test/java/org/apache/ignite/internal/thread/context/ContextAttributesTest.java

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333

3434
import static java.util.concurrent.TimeUnit.MILLISECONDS;
3535
import static org.apache.ignite.testframework.GridTestUtils.assertThrowsAnyCause;
36+
import static org.apache.ignite.testframework.GridTestUtils.assertThrowsWithCause;
3637

3738
/** */
3839
public class ContextAttributesTest extends GridCommonAbstractTest {
@@ -273,12 +274,32 @@ public void testMaximumAttributesInstanceCount() {
273274
);
274275
}
275276

277+
/** */
278+
@Test
279+
public void testUnorderedScopeClosing() {
280+
Scope scope1 = Context.Builder.create().with(INT_ATTR, 0).build().attach();
281+
282+
try {
283+
try (Scope ignored = Context.Builder.create().with(STR_ATTR, "test").build().attach()) {
284+
assertThrowsWithCause(scope1::close, AssertionError.class);
285+
}
286+
}
287+
finally {
288+
scope1.close();
289+
}
290+
291+
checkAttributeValues(DFLT_STR_VAL, DFLT_INT_VAL);
292+
293+
assertThrowsWithCause(scope1::close, AssertionError.class);
294+
}
295+
296+
276297
/** */
277298
@Test
278299
public void testEmptySnapshot() {
279300
checkAttributeValues(DFLT_STR_VAL, DFLT_INT_VAL);
280301

281-
ContextSnapshot snapshot = ContextSnapshot.capture();
302+
ContextSnapshot snapshot = Context.createSnapshot();
282303

283304
try (Scope ignored = snapshot.restore()) {
284305
checkAttributeValues(DFLT_STR_VAL, DFLT_INT_VAL);
@@ -297,7 +318,7 @@ public void testSnapshot() {
297318
try (Scope ignored = Context.Builder.create().with(INT_ATTR, 1).with(STR_ATTR, "test1").build().attach()) {
298319
checkAttributeValues("test1", 1);
299320

300-
snapshot = ContextSnapshot.capture();
321+
snapshot = Context.createSnapshot();
301322
}
302323

303324
checkAttributeValues(DFLT_STR_VAL, DFLT_INT_VAL);
@@ -318,7 +339,7 @@ public void testNestedScopeSnapshot() {
318339
try (Scope ignored2 = Context.Builder.create().with(STR_ATTR, "test2").build().attach()) {
319340
checkAttributeValues("test2", 1);
320341

321-
snapshot = ContextSnapshot.capture();
342+
snapshot = Context.createSnapshot();
322343
}
323344
}
324345

@@ -337,7 +358,7 @@ public void testNestedScopeInSnapshotScope() {
337358
try (Scope ignored = Context.Builder.create().with(INT_ATTR, 1).with(STR_ATTR, "test1").build().attach()) {
338359
checkAttributeValues("test1", 1);
339360

340-
snapshot0 = ContextSnapshot.capture();
361+
snapshot0 = Context.createSnapshot();
341362
}
342363

343364
ContextSnapshot snapshot1;
@@ -348,7 +369,7 @@ public void testNestedScopeInSnapshotScope() {
348369
try (Scope ignored2 = Context.Builder.create().with(INT_ATTR, 2).build().attach()) {
349370
checkAttributeValues("test1", 2);
350371

351-
snapshot1 = ContextSnapshot.capture();
372+
snapshot1 = Context.createSnapshot();
352373
}
353374

354375
checkAttributeValues("test1", 1);
@@ -369,7 +390,7 @@ public void testSnapshotRestoreInExistingScope() {
369390
try (Scope ignored = Context.Builder.create().with(STR_ATTR, "test1").build().attach()) {
370391
checkAttributeValues("test1", DFLT_INT_VAL);
371392

372-
snapshot = ContextSnapshot.capture();
393+
snapshot = Context.createSnapshot();
373394
}
374395

375396
try (Scope ignored1 = Context.Builder.create().with(INT_ATTR, 1).build().attach()) {
@@ -386,6 +407,28 @@ public void testSnapshotRestoreInExistingScope() {
386407
checkAttributeValues(DFLT_STR_VAL, DFLT_INT_VAL);
387408
}
388409

410+
/** */
411+
@Test
412+
public void testSnapshotScopeUnorderedClosing() {
413+
ContextSnapshot snapshot;
414+
415+
try (Scope ignored = Context.Builder.create().with(STR_ATTR, "test1").build().attach()) {
416+
checkAttributeValues("test1", DFLT_INT_VAL);
417+
418+
snapshot = Context.createSnapshot();
419+
}
420+
421+
try (Scope snpScope = snapshot.restore()) {
422+
try (Scope ignored1 = Context.Builder.create().with(INT_ATTR, 2).build().attach()) {
423+
checkAttributeValues("test1", 2);
424+
425+
assertThrowsWithCause(snpScope::close, AssertionError.class);
426+
}
427+
}
428+
429+
checkAttributeValues(DFLT_STR_VAL, DFLT_INT_VAL);
430+
}
431+
389432
/** */
390433
@Test
391434
public void testScopedContextAwareThreadPool() throws Exception {

0 commit comments

Comments
 (0)