Skip to content

Commit fe13e4b

Browse files
mccullsmcculls
authored
Simplify management of active message scopes by tracking them per-session-thread (DataDog#3027)
This assumes that if a thread receives a message from the same session, any previous message scope associated with that thread is closed whether the message is received via the same consumer or not. All active message scopes are still closed when the session is closed. We also clean-up stale message scopes from JMS sessions if the associated threads have stopped.
1 parent d6e9d05 commit fe13e4b

File tree

4 files changed

+62
-36
lines changed

4 files changed

+62
-36
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/jms/MessageConsumerState.java‎

Lines changed: 3 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,9 @@
22

33
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
44
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
5-
import java.util.Map;
6-
import java.util.concurrent.ConcurrentHashMap;
75

86
/** Tracks message scopes and spans when consuming messages with {@code receive}. */
97
public final class MessageConsumerState {
10-
private final Map<Thread, AgentScope> currentScopes = new ConcurrentHashMap<>();
118

129
private final SessionState sessionState;
1310
private final UTF8BytesString resourceName;
@@ -18,8 +15,6 @@ public MessageConsumerState(
1815
this.sessionState = sessionState;
1916
this.resourceName = resourceName;
2017
this.propagationDisabled = propagationDisabled;
21-
22-
sessionState.registerConsumerState(this);
2318
}
2419

2520
public SessionState getSessionState() {
@@ -36,27 +31,11 @@ public boolean isPropagationDisabled() {
3631

3732
/** Closes the given message scope when the next message is consumed or the consumer is closed. */
3833
public void closeOnIteration(AgentScope newScope) {
39-
maybeCloseScope(currentScopes.put(Thread.currentThread(), newScope));
34+
sessionState.closeOnIteration(newScope); // tracked per-session-thread
4035
}
4136

37+
/** Closes the scope previously registered by closeOnIteration, assumes same calling thread. */
4238
public void closePreviousMessageScope() {
43-
maybeCloseScope(currentScopes.remove(Thread.currentThread()));
44-
}
45-
46-
public void onClose() {
47-
for (AgentScope scope : currentScopes.values()) {
48-
maybeCloseScope(scope);
49-
}
50-
currentScopes.clear();
51-
sessionState.unregisterConsumerState(this);
52-
}
53-
54-
private void maybeCloseScope(AgentScope scope) {
55-
if (null != scope) {
56-
scope.close();
57-
if (sessionState.isAutoAcknowledge()) {
58-
scope.span().finish();
59-
}
60-
}
39+
sessionState.closePreviousMessageScope(); // tracked per-session-thread
6140
}
6241
}

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/jms/SessionState.java‎

Lines changed: 42 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
package datadog.trace.bootstrap.instrumentation.jms;
22

3-
import static java.util.Collections.newSetFromMap;
4-
3+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
54
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
5+
import java.util.Iterator;
6+
import java.util.Map;
67
import java.util.Queue;
7-
import java.util.Set;
88
import java.util.concurrent.ArrayBlockingQueue;
99
import java.util.concurrent.ConcurrentHashMap;
1010
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
@@ -16,6 +16,9 @@
1616
*/
1717
public final class SessionState {
1818

19+
private static final AtomicIntegerFieldUpdater<SessionState> SCOPE_COUNT =
20+
AtomicIntegerFieldUpdater.newUpdater(SessionState.class, "scopeCount");
21+
1922
private static final AtomicReferenceFieldUpdater<SessionState, Queue> CAPTURED_SPANS =
2023
AtomicReferenceFieldUpdater.newUpdater(SessionState.class, Queue.class, "capturedSpans");
2124
private static final AtomicIntegerFieldUpdater<SessionState> SPAN_COUNT =
@@ -28,8 +31,8 @@ public final class SessionState {
2831
private final int ackMode;
2932

3033
// consumer-related session state
31-
private final Set<MessageConsumerState> consumerStates =
32-
newSetFromMap(new ConcurrentHashMap<MessageConsumerState, Boolean>());
34+
private final Map<Thread, AgentScope> activeScopes = new ConcurrentHashMap<>();
35+
private volatile int scopeCount;
3336
private volatile Queue<AgentSpan> capturedSpans;
3437
private volatile int spanCount;
3538

@@ -93,20 +96,48 @@ public void onCommitOrRollback() {
9396
}
9497
}
9598

96-
void registerConsumerState(MessageConsumerState consumerState) {
97-
consumerStates.add(consumerState);
99+
/** Closes the given message scope when the next message is consumed or the session is closed. */
100+
void closeOnIteration(AgentScope newScope) {
101+
if (SCOPE_COUNT.incrementAndGet(this) > 100) {
102+
closeStaleScopes();
103+
}
104+
maybeCloseScope(activeScopes.put(Thread.currentThread(), newScope));
98105
}
99106

100-
void unregisterConsumerState(MessageConsumerState consumerState) {
101-
consumerStates.remove(consumerState);
107+
/** Closes the scope previously registered by closeOnIteration, assumes same calling thread. */
108+
void closePreviousMessageScope() {
109+
maybeCloseScope(activeScopes.remove(Thread.currentThread()));
102110
}
103111

112+
/** Closes any active message scopes and finishes any pending transacted spans. */
104113
public void onClose() {
105-
for (MessageConsumerState consumerState : consumerStates) {
106-
consumerState.onClose(); // eventually calls unregisterConsumerState
114+
for (AgentScope scope : activeScopes.values()) {
115+
maybeCloseScope(scope);
107116
}
117+
activeScopes.clear();
108118
if (isTransactedSession()) {
109119
onCommitOrRollback(); // implicit rollback of any active transaction
110120
}
111121
}
122+
123+
private void maybeCloseScope(AgentScope scope) {
124+
if (null != scope) {
125+
SCOPE_COUNT.decrementAndGet(this);
126+
scope.close();
127+
if (isAutoAcknowledge()) {
128+
scope.span().finish();
129+
}
130+
}
131+
}
132+
133+
private void closeStaleScopes() {
134+
Iterator<Map.Entry<Thread, AgentScope>> itr = activeScopes.entrySet().iterator();
135+
while (itr.hasNext()) {
136+
Map.Entry<Thread, AgentScope> entry = itr.next();
137+
if (!entry.getKey().isAlive()) {
138+
maybeCloseScope(entry.getValue());
139+
itr.remove();
140+
}
141+
}
142+
}
112143
}

‎dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/jms/SessionStateTest.groovy‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.bootstrap.instrumentation.jms
22

3+
import datadog.trace.bootstrap.instrumentation.api.AgentScope
34
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
45
import datadog.trace.test.util.DDSpecification
56

@@ -49,4 +50,19 @@ class SessionStateTest extends DDSpecification {
4950
0 * span2.finish()
5051
sessionState.capturedSpanCount == 1
5152
}
53+
54+
def "stale scopes are cleaned up from session"() {
55+
setup:
56+
SessionState sessionState = new SessionState(0)
57+
when: "add thread scopes without triggering cleanup"
58+
for (int i = 0; i < 100; ++i) {
59+
Thread.start { sessionState.closeOnIteration(Mock(AgentScope)) }.join()
60+
}
61+
then:
62+
sessionState.scopeCount == 100
63+
when: "trigger cleanup"
64+
sessionState.closeOnIteration(Mock(AgentScope))
65+
then:
66+
sessionState.scopeCount == 1
67+
}
5268
}

‎dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ public static void beforeClose(@Advice.This final MessageConsumer consumer) {
145145
InstrumentationContext.get(MessageConsumer.class, MessageConsumerState.class)
146146
.get(consumer);
147147
if (null != consumerState) {
148-
consumerState.onClose();
148+
consumerState.closePreviousMessageScope();
149149
}
150150
}
151151
}

0 commit comments

Comments
 (0)