Skip to content

Commit 3d19b84

Browse files
diemoldiemol
committed
[grid] Draining a Node after X sessions have been created.
This will help to enable a Dynamic Grid in Kubernetes, as one can create a container with a single session which will shut down on its own after the session is completed. Helps with #9845 and SeleniumHQ/docker-selenium#1514
1 parent ecc30e4 commit 3d19b84

5 files changed

Lines changed: 106 additions & 26 deletions

File tree

‎java/src/org/openqa/selenium/grid/node/config/NodeFlags.java‎

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -17,19 +17,6 @@
1717

1818
package org.openqa.selenium.grid.node.config;
1919

20-
import static org.openqa.selenium.grid.config.StandardGridRoles.NODE_ROLE;
21-
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_DETECT_DRIVERS;
22-
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_HEARTBEAT_PERIOD;
23-
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_MAX_SESSIONS;
24-
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_NODE_IMPLEMENTATION;
25-
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_NO_VNC_PORT;
26-
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_REGISTER_CYCLE;
27-
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_REGISTER_PERIOD;
28-
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_SESSION_TIMEOUT;
29-
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_VNC_ENV_VAR;
30-
import static org.openqa.selenium.grid.node.config.NodeOptions.NODE_SECTION;
31-
import static org.openqa.selenium.grid.node.config.NodeOptions.OVERRIDE_MAX_SESSIONS;
32-
3320
import com.google.auto.service.AutoService;
3421

3522
import com.beust.jcommander.Parameter;
@@ -44,6 +31,20 @@
4431
import java.util.List;
4532
import java.util.Set;
4633

34+
import static org.openqa.selenium.grid.config.StandardGridRoles.NODE_ROLE;
35+
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_DETECT_DRIVERS;
36+
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_DRAIN_AFTER_SESSION_COUNT;
37+
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_HEARTBEAT_PERIOD;
38+
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_MAX_SESSIONS;
39+
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_NODE_IMPLEMENTATION;
40+
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_NO_VNC_PORT;
41+
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_REGISTER_CYCLE;
42+
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_REGISTER_PERIOD;
43+
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_SESSION_TIMEOUT;
44+
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_VNC_ENV_VAR;
45+
import static org.openqa.selenium.grid.node.config.NodeOptions.NODE_SECTION;
46+
import static org.openqa.selenium.grid.node.config.NodeOptions.OVERRIDE_MAX_SESSIONS;
47+
4748
@SuppressWarnings({"unused", "FieldMayBeFinal"})
4849
@AutoService(HasRoles.class)
4950
public class NodeFlags implements HasRoles {
@@ -187,6 +188,13 @@ public class NodeFlags implements HasRoles {
187188
@ConfigValue(section = NODE_SECTION, name = "no-vnc-port", example = "7900")
188189
public int noVncPort = DEFAULT_NO_VNC_PORT;
189190

191+
@Parameter(
192+
names = "--drain-after-session-count",
193+
description = "Drain and shutdown the Node after X sessions have been executed. Useful for " +
194+
"environments like Kubernetes. A value higher than zero enables this feature.")
195+
@ConfigValue(section = NODE_SECTION, name = "drain-after-session-count", example = "1")
196+
public int drainAfterSessionCount = DEFAULT_DRAIN_AFTER_SESSION_COUNT;
197+
190198
@Parameter(
191199
names = {"--node-implementation"},
192200
description = "Full classname of non-default Node implementation. This is used to manage "

‎java/src/org/openqa/selenium/grid/node/config/NodeOptions.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ public class NodeOptions {
6969
public static final int DEFAULT_MAX_SESSIONS = Runtime.getRuntime().availableProcessors();
7070
public static final int DEFAULT_HEARTBEAT_PERIOD = 60;
7171
public static final int DEFAULT_SESSION_TIMEOUT = 300;
72+
public static final int DEFAULT_DRAIN_AFTER_SESSION_COUNT = 0;
7273
static final String NODE_SECTION = "node";
7374
static final boolean DEFAULT_DETECT_DRIVERS = true;
7475
static final boolean OVERRIDE_MAX_SESSIONS = false;
@@ -232,6 +233,13 @@ public Duration getSessionTimeout() {
232233
return Duration.ofSeconds(seconds);
233234
}
234235

236+
public int getDrainAfterSessionCount() {
237+
return Math.max(
238+
config.getInt(NODE_SECTION, "drain-after-session-count")
239+
.orElse(DEFAULT_DRAIN_AFTER_SESSION_COUNT),
240+
DEFAULT_DRAIN_AFTER_SESSION_COUNT);
241+
}
242+
235243
@VisibleForTesting
236244
boolean isVncEnabled() {
237245
String vncEnvVar = config.get(NODE_SECTION, "vnc-env-var").orElse(DEFAULT_VNC_ENV_VAR);

‎java/src/org/openqa/selenium/grid/node/local/LocalNode.java‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@
8888
import java.util.concurrent.Executors;
8989
import java.util.concurrent.ScheduledExecutorService;
9090
import java.util.concurrent.TimeUnit;
91+
import java.util.concurrent.atomic.AtomicBoolean;
9192
import java.util.concurrent.atomic.AtomicInteger;
9293
import java.util.logging.Logger;
9394
import java.util.stream.Collectors;
@@ -116,10 +117,13 @@ public class LocalNode extends Node {
116117
private final Duration heartbeatPeriod;
117118
private final HealthCheck healthCheck;
118119
private final int maxSessionCount;
120+
private final int configuredSessionCount;
121+
private final AtomicBoolean drainAfterSessions = new AtomicBoolean();
119122
private final List<SessionSlot> factories;
120123
private final Cache<SessionId, SessionSlot> currentSessions;
121124
private final Cache<SessionId, TemporaryFilesystem> tempFileSystems;
122125
private final AtomicInteger pendingSessions = new AtomicInteger();
126+
private final AtomicInteger sessionCount = new AtomicInteger();
123127

124128
private LocalNode(
125129
Tracer tracer,
@@ -128,6 +132,7 @@ private LocalNode(
128132
URI gridUri,
129133
HealthCheck healthCheck,
130134
int maxSessionCount,
135+
int drainAfterSessionCount,
131136
Ticker ticker,
132137
Duration sessionTimeout,
133138
Duration heartbeatPeriod,
@@ -139,10 +144,14 @@ private LocalNode(
139144

140145
this.externalUri = Require.nonNull("Remote node URI", uri);
141146
this.gridUri = Require.nonNull("Grid URI", gridUri);
142-
this.maxSessionCount = Math.min(Require.positive("Max session count", maxSessionCount), factories.size());
147+
this.maxSessionCount = Math.min(
148+
Require.positive("Max session count", maxSessionCount), factories.size());
143149
this.heartbeatPeriod = heartbeatPeriod;
144150
this.factories = ImmutableList.copyOf(factories);
145151
Require.nonNull("Registration secret", registrationSecret);
152+
this.configuredSessionCount = drainAfterSessionCount;
153+
this.drainAfterSessions.set(this.configuredSessionCount > 0);
154+
this.sessionCount.set(drainAfterSessionCount);
146155

147156
this.healthCheck = healthCheck == null ?
148157
() -> {
@@ -345,6 +354,8 @@ public Either<WebDriverException, CreateSessionResponse> newSession(CreateSessio
345354
ActiveSession session = possibleSession.right();
346355
currentSessions.put(session.getId(), slotToUse);
347356

357+
checkSessionCount();
358+
348359
SessionId sessionId = session.getId();
349360
Capabilities caps = session.getCapabilities();
350361
SESSION_ID.accept(span, sessionId);
@@ -597,6 +608,20 @@ public void drain() {
597608
}
598609
}
599610

611+
private void checkSessionCount() {
612+
if (this.drainAfterSessions.get()) {
613+
int remainingSessions = this.sessionCount.decrementAndGet();
614+
LOG.log(
615+
Debug.getDebugLogLevel(),
616+
String.format("%s remaining sessions before draining Node", remainingSessions));
617+
if (remainingSessions <= 0) {
618+
LOG.info(String.format("Draining Node, configured sessions value (%s) has been reached.",
619+
this.configuredSessionCount));
620+
drain();
621+
}
622+
}
623+
}
624+
600625
private Map<String, Object> toJson() {
601626
return ImmutableMap.of(
602627
"id", getId(),
@@ -616,7 +641,8 @@ public static class Builder {
616641
private final URI gridUri;
617642
private final Secret registrationSecret;
618643
private final ImmutableList.Builder<SessionSlot> factories;
619-
private int maxCount = NodeOptions.DEFAULT_MAX_SESSIONS;
644+
private int maxSessions = NodeOptions.DEFAULT_MAX_SESSIONS;
645+
private int drainAfterSessionCount = NodeOptions.DEFAULT_DRAIN_AFTER_SESSION_COUNT;
620646
private Ticker ticker = Ticker.systemTicker();
621647
private Duration sessionTimeout = Duration.ofSeconds(NodeOptions.DEFAULT_SESSION_TIMEOUT);
622648
private HealthCheck healthCheck;
@@ -646,7 +672,12 @@ public Builder add(Capabilities stereotype, SessionFactory factory) {
646672
}
647673

648674
public Builder maximumConcurrentSessions(int maxCount) {
649-
this.maxCount = Require.positive("Max session count", maxCount);
675+
this.maxSessions = Require.positive("Max session count", maxCount);
676+
return this;
677+
}
678+
679+
public Builder drainAfterSessionCount(int sessionCount) {
680+
this.drainAfterSessionCount = sessionCount;
650681
return this;
651682
}
652683

@@ -667,7 +698,8 @@ public LocalNode build() {
667698
uri,
668699
gridUri,
669700
healthCheck,
670-
maxCount,
701+
maxSessions,
702+
drainAfterSessionCount,
671703
ticker,
672704
sessionTimeout,
673705
heartbeatPeriod,

‎java/src/org/openqa/selenium/grid/node/local/LocalNodeFactory.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,9 +68,9 @@ public static Node create(Config config) {
6868
secretOptions.getRegistrationSecret())
6969
.maximumConcurrentSessions(nodeOptions.getMaxSessions())
7070
.sessionTimeout(sessionTimeout)
71+
.drainAfterSessionCount(nodeOptions.getDrainAfterSessionCount())
7172
.heartbeatPeriod(nodeOptions.getHeartbeatPeriod());
7273

73-
7474
List<DriverService.Builder<?, ?>> builders = new ArrayList<>();
7575
ServiceLoader.load(DriverService.Builder.class).forEach(builders::add);
7676

‎java/test/org/openqa/selenium/grid/node/local/LocalNodeTest.java‎

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import com.google.common.collect.ImmutableMap;
2121
import com.google.common.collect.ImmutableSet;
22+
2223
import org.junit.Before;
2324
import org.junit.Test;
2425
import org.openqa.selenium.Capabilities;
@@ -51,6 +52,7 @@
5152
import java.time.Instant;
5253
import java.util.ArrayList;
5354
import java.util.List;
55+
import java.util.Objects;
5456
import java.util.Optional;
5557
import java.util.concurrent.Callable;
5658
import java.util.concurrent.Executors;
@@ -160,30 +162,30 @@ public void cannotCreateNewSessionsOnMaxSessionCount() {
160162
public void canReturnStatusInfo() {
161163
NodeStatus status = node.getStatus();
162164
assertThat(status.getSlots().stream()
163-
.filter(slot -> slot.getSession()!=null)
164-
.map(Slot::getSession)
165+
.map(Slot::getSession)
166+
.filter(Objects::nonNull)
165167
.filter(s -> s.getId().equals(session.getId()))).isNotEmpty();
166168

167169
node.stop(session.getId());
168170
status = node.getStatus();
169171
assertThat(status.getSlots().stream()
170-
.filter(slot -> slot.getSession()!=null)
171-
.map(Slot::getSession)
172+
.map(Slot::getSession)
173+
.filter(Objects::nonNull)
172174
.filter(s -> s.getId().equals(session.getId()))).isEmpty();
173175
}
174176

175177
@Test
176178
public void nodeStatusInfoIsImmutable() {
177179
NodeStatus status = node.getStatus();
178180
assertThat(status.getSlots().stream()
179-
.filter(slot -> slot.getSession()!=null)
180-
.map(slot -> slot.getSession())
181+
.map(Slot::getSession)
182+
.filter(Objects::nonNull)
181183
.filter(s -> s.getId().equals(session.getId()))).isNotEmpty();
182184

183185
node.stop(session.getId());
184186
assertThat(status.getSlots().stream()
185-
.filter(slot -> slot.getSession()!=null)
186-
.map(slot -> slot.getSession())
187+
.map(Slot::getSession)
188+
.filter(Objects::nonNull)
187189
.filter(s -> s.getId().equals(session.getId()))).isNotEmpty();
188190
}
189191

@@ -242,4 +244,34 @@ public HttpResponse execute(HttpRequest req) {
242244
assertThat(res.isSuccessful()).isTrue();
243245
}
244246
}
247+
248+
@Test
249+
public void nodeDrainsAfterSessionCountIsReached() throws URISyntaxException {
250+
Tracer tracer = DefaultTestTracer.createTracer();
251+
EventBus bus = new GuavaEventBus();
252+
URI uri = new URI("http://localhost:5678");
253+
Capabilities stereotype = new ImmutableCapabilities("browserName", "bread");
254+
255+
LocalNode.Builder builder = LocalNode.builder(tracer, bus, uri, uri, registrationSecret)
256+
.maximumConcurrentSessions(10)
257+
.drainAfterSessionCount(5);
258+
for (int i = 0; i < 5; i++) {
259+
builder.add(stereotype, new TestSessionFactory(
260+
(id, caps) -> new Session(id, uri, stereotype, caps, Instant.now())));
261+
}
262+
LocalNode localNode = builder.build();
263+
264+
assertThat(localNode.isDraining()).isFalse();
265+
266+
for (int i = 0; i < 5; i++) {
267+
Either<WebDriverException, CreateSessionResponse> response = localNode.newSession(
268+
new CreateSessionRequest(
269+
ImmutableSet.of(W3C),
270+
stereotype,
271+
ImmutableMap.of()));
272+
assertThat(response.isRight()).isTrue();
273+
}
274+
275+
assertThat(localNode.isDraining()).isTrue();
276+
}
245277
}

0 commit comments

Comments
 (0)