Skip to content

Commit 0f831ba

Browse files
VietND96VietND96
authored
[grid] Fix race condition and improve logging in LocalSessionMap (#15370)
* [grid] Fix race condition in LocalSessionMap and improve logging * Add LocalSessionMapTest * Add suggestion to improve code --------- Signed-off-by: Viet Nguyen Duc <[email protected]>
1 parent 8e46f06 commit 0f831ba

File tree

3 files changed

+882
-24
lines changed

3 files changed

+882
-24
lines changed

‎java/src/org/openqa/selenium/grid/sessionmap/local/LocalSessionMap.java

Lines changed: 160 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,18 @@
2020
import static org.openqa.selenium.remote.RemoteTags.SESSION_ID;
2121
import static org.openqa.selenium.remote.RemoteTags.SESSION_ID_EVENT;
2222

23-
import java.util.List;
23+
import java.net.URI;
24+
import java.util.Collection;
25+
import java.util.Collections;
26+
import java.util.HashMap;
27+
import java.util.HashSet;
2428
import java.util.Map;
29+
import java.util.Set;
2530
import java.util.concurrent.ConcurrentHashMap;
2631
import java.util.concurrent.ConcurrentMap;
2732
import java.util.logging.Logger;
28-
import java.util.stream.Collectors;
2933
import org.openqa.selenium.NoSuchSessionException;
34+
import org.openqa.selenium.events.Event;
3035
import org.openqa.selenium.events.EventBus;
3136
import org.openqa.selenium.grid.config.Config;
3237
import org.openqa.selenium.grid.data.NodeRemovedEvent;
@@ -48,7 +53,7 @@ public class LocalSessionMap extends SessionMap {
4853
private static final Logger LOG = Logger.getLogger(LocalSessionMap.class.getName());
4954

5055
private final EventBus bus;
51-
private final ConcurrentMap<SessionId, Session> knownSessions = new ConcurrentHashMap<>();
56+
private final IndexedSessionMap knownSessions = new IndexedSessionMap();
5257

5358
public LocalSessionMap(Tracer tracer, EventBus bus) {
5459
super(tracer);
@@ -59,23 +64,14 @@ public LocalSessionMap(Tracer tracer, EventBus bus) {
5964

6065
bus.addListener(
6166
NodeRemovedEvent.listener(
62-
nodeStatus ->
63-
nodeStatus.getSlots().stream()
64-
.filter(slot -> slot.getSession() != null)
65-
.map(slot -> slot.getSession().getId())
66-
.forEach(this::remove)));
67+
nodeStatus -> {
68+
batchRemoveByUri(nodeStatus.getExternalUri(), NodeRemovedEvent.class);
69+
}));
6770

6871
bus.addListener(
6972
NodeRestartedEvent.listener(
7073
previousNodeStatus -> {
71-
List<SessionId> toRemove =
72-
knownSessions.entrySet().stream()
73-
.filter(
74-
(e) -> e.getValue().getUri().equals(previousNodeStatus.getExternalUri()))
75-
.map(Map.Entry::getKey)
76-
.collect(Collectors.toList());
77-
78-
toRemove.forEach(this::remove);
74+
batchRemoveByUri(previousNodeStatus.getExternalUri(), NodeRestartedEvent.class);
7975
}));
8076
}
8177

@@ -95,17 +91,23 @@ public boolean isReady() {
9591
public boolean add(Session session) {
9692
Require.nonNull("Session", session);
9793

94+
SessionId id = session.getId();
95+
knownSessions.put(id, session);
96+
9897
try (Span span = tracer.getCurrentContext().createSpan("local_sessionmap.add")) {
9998
AttributeMap attributeMap = tracer.createAttributeMap();
10099
attributeMap.put(AttributeKey.LOGGER_CLASS.getKey(), getClass().getName());
101-
SessionId id = session.getId();
102100
SESSION_ID.accept(span, id);
103101
SESSION_ID_EVENT.accept(attributeMap, id);
104-
knownSessions.put(session.getId(), session);
105-
span.addEvent("Added session into local session map", attributeMap);
106102

107-
return true;
103+
String sessionAddedMessage =
104+
String.format(
105+
"Added session to local Session Map, Id: %s, Node: %s", id, session.getUri());
106+
span.addEvent(sessionAddedMessage, attributeMap);
107+
LOG.info(sessionAddedMessage);
108108
}
109+
110+
return true;
109111
}
110112

111113
@Override
@@ -116,23 +118,157 @@ public Session get(SessionId id) {
116118
if (session == null) {
117119
throw new NoSuchSessionException("Unable to find session with ID: " + id);
118120
}
119-
120121
return session;
121122
}
122123

123124
@Override
124125
public void remove(SessionId id) {
125126
Require.nonNull("Session ID", id);
126127

128+
Session removedSession = knownSessions.remove(id);
129+
127130
try (Span span = tracer.getCurrentContext().createSpan("local_sessionmap.remove")) {
128131
AttributeMap attributeMap = tracer.createAttributeMap();
129132
attributeMap.put(AttributeKey.LOGGER_CLASS.getKey(), getClass().getName());
130133
SESSION_ID.accept(span, id);
131134
SESSION_ID_EVENT.accept(attributeMap, id);
132-
knownSessions.remove(id);
133-
String sessionDeletedMessage = "Deleted session from local Session Map";
135+
136+
String sessionDeletedMessage =
137+
String.format(
138+
"Deleted session from local Session Map, Id: %s, Node: %s",
139+
id,
140+
removedSession != null ? String.valueOf(removedSession.getUri()) : "unidentified");
134141
span.addEvent(sessionDeletedMessage, attributeMap);
135-
LOG.info(String.format("%s, Id: %s", sessionDeletedMessage, id));
142+
LOG.info(sessionDeletedMessage);
143+
}
144+
}
145+
146+
private void batchRemoveByUri(URI externalUri, Class<? extends Event> eventClass) {
147+
Set<SessionId> sessionsToRemove = knownSessions.getSessionsByUri(externalUri);
148+
149+
if (sessionsToRemove.isEmpty()) {
150+
return; // Early return for empty operations - no tracing overhead
151+
}
152+
153+
knownSessions.batchRemove(sessionsToRemove);
154+
155+
try (Span span = tracer.getCurrentContext().createSpan("local_sessionmap.batch_remove")) {
156+
AttributeMap attributeMap = tracer.createAttributeMap();
157+
attributeMap.put(AttributeKey.LOGGER_CLASS.getKey(), getClass().getName());
158+
attributeMap.put("event.class", eventClass.getName());
159+
attributeMap.put("node.uri", externalUri.toString());
160+
attributeMap.put("sessions.count", sessionsToRemove.size());
161+
162+
String batchRemoveMessage =
163+
String.format(
164+
"Batch removed %d sessions from local Session Map for Node %s (triggered by %s)",
165+
sessionsToRemove.size(), externalUri, eventClass.getSimpleName());
166+
span.addEvent(batchRemoveMessage, attributeMap);
167+
LOG.info(batchRemoveMessage);
168+
}
169+
}
170+
171+
private static class IndexedSessionMap {
172+
private final ConcurrentMap<SessionId, Session> sessions = new ConcurrentHashMap<>();
173+
private final ConcurrentMap<URI, Set<SessionId>> sessionsByUri = new ConcurrentHashMap<>();
174+
private final Object coordinationLock = new Object();
175+
176+
public Session get(SessionId id) {
177+
return sessions.get(id);
178+
}
179+
180+
public Session put(SessionId id, Session session) {
181+
synchronized (coordinationLock) {
182+
Session previous = sessions.put(id, session);
183+
184+
if (previous != null && previous.getUri() != null) {
185+
cleanupUriIndex(previous.getUri(), id);
186+
}
187+
188+
URI sessionUri = session.getUri();
189+
if (sessionUri != null) {
190+
sessionsByUri.computeIfAbsent(sessionUri, k -> ConcurrentHashMap.newKeySet()).add(id);
191+
}
192+
193+
return previous;
194+
}
195+
}
196+
197+
public Session remove(SessionId id) {
198+
synchronized (coordinationLock) {
199+
Session removed = sessions.remove(id);
200+
201+
if (removed != null && removed.getUri() != null) {
202+
cleanupUriIndex(removed.getUri(), id);
203+
}
204+
205+
return removed;
206+
}
207+
}
208+
209+
public void batchRemove(Set<SessionId> sessionIds) {
210+
synchronized (coordinationLock) {
211+
Map<URI, Set<SessionId>> uriToSessionIds = new HashMap<>();
212+
213+
// Single loop: remove sessions and collect URI mappings in one pass
214+
for (SessionId id : sessionIds) {
215+
Session session = sessions.remove(id);
216+
if (session != null && session.getUri() != null) {
217+
uriToSessionIds.computeIfAbsent(session.getUri(), k -> new HashSet<>()).add(id);
218+
}
219+
}
220+
221+
// Clean up URI index for all affected URIs
222+
for (Map.Entry<URI, Set<SessionId>> entry : uriToSessionIds.entrySet()) {
223+
cleanupUriIndex(entry.getKey(), entry.getValue());
224+
}
225+
}
226+
}
227+
228+
private void cleanupUriIndex(URI uri, SessionId sessionId) {
229+
sessionsByUri.computeIfPresent(
230+
uri,
231+
(key, sessionIds) -> {
232+
sessionIds.remove(sessionId);
233+
return sessionIds.isEmpty() ? null : sessionIds;
234+
});
235+
}
236+
237+
private void cleanupUriIndex(URI uri, Set<SessionId> sessionIdsToRemove) {
238+
sessionsByUri.computeIfPresent(
239+
uri,
240+
(key, sessionIds) -> {
241+
sessionIds.removeAll(sessionIdsToRemove);
242+
return sessionIds.isEmpty() ? null : sessionIds;
243+
});
244+
}
245+
246+
public Set<SessionId> getSessionsByUri(URI uri) {
247+
Set<SessionId> result = sessionsByUri.get(uri);
248+
return (result != null && !result.isEmpty()) ? result : Set.of();
249+
}
250+
251+
public Set<Map.Entry<SessionId, Session>> entrySet() {
252+
return Collections.unmodifiableSet(sessions.entrySet());
253+
}
254+
255+
public Collection<Session> values() {
256+
return Collections.unmodifiableCollection(sessions.values());
257+
}
258+
259+
public int size() {
260+
return sessions.size();
261+
}
262+
263+
public boolean isEmpty() {
264+
return sessions.isEmpty();
265+
}
266+
267+
public void clear() {
268+
synchronized (coordinationLock) {
269+
sessions.clear();
270+
sessionsByUri.clear();
271+
}
136272
}
137273
}
138274
}

‎java/test/org/openqa/selenium/grid/sessionmap/local/BUILD.bazel

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
load("@rules_jvm_external//:defs.bzl", "artifact")
2+
load("//java:defs.bzl", "JUNIT5_DEPS", "java_test_suite")
3+
4+
java_test_suite(
5+
name = "SmallTests",
6+
size = "medium",
7+
srcs = glob(["*.java"]),
8+
deps = [
9+
"//java/src/org/openqa/selenium:core",
10+
"//java/src/org/openqa/selenium/events",
11+
"//java/src/org/openqa/selenium/events/local",
12+
"//java/src/org/openqa/selenium/grid/data",
13+
"//java/src/org/openqa/selenium/grid/sessionmap",
14+
"//java/src/org/openqa/selenium/grid/sessionmap/local",
15+
"//java/src/org/openqa/selenium/remote",
16+
"//java/test/org/openqa/selenium/remote/tracing:tracing-support",
17+
artifact("org.assertj:assertj-core"),
18+
artifact("org.junit.jupiter:junit-jupiter-api"),
19+
] + JUNIT5_DEPS,
20+
)

0 commit comments

Comments
 (0)