[grid]: Add a new session to SessionMap in Distributor

This has one major advantage: anything that extends Node
now doesn't need to be aware of the SessionMap, though
it still needs to know enough to fire the close event
when the session is finally closed.

java/server/src/org/openqa/selenium/grid/commands/Hub.java

@@ -114,7 +114,8 @@ public Executable configure(String... args) {
       Distributor distributor = new LocalDistributor(
           tracer,
           bus,
-          clientFactory);
+          clientFactory,
+          sessions);
       Router router = new Router(tracer, clientFactory, sessions, distributor);
 
       Server<?> server = new BaseServer<>(

java/server/src/org/openqa/selenium/grid/commands/Standalone.java

@@ -119,7 +119,7 @@ public Executable configure(String... args) {
       HttpClient.Factory clientFactory = HttpClient.Factory.createDefault();
 
       SessionMap sessions = new LocalSessionMap(tracer, bus);
-      Distributor distributor = new LocalDistributor(tracer, bus, clientFactory);
+      Distributor distributor = new LocalDistributor(tracer, bus, clientFactory, sessions);
       Router router = new Router(tracer, clientFactory, sessions, distributor);
 
       String hostName;
@@ -142,8 +142,7 @@ public Executable configure(String... args) {
           tracer,
           bus,
           clientFactory,
-          localhost,
-          sessions)
+          localhost)
           .maximumConcurrentSessions(Runtime.getRuntime().availableProcessors() * 3);
       nodeFlags.configure(config, clientFactory, node);
 

java/server/src/org/openqa/selenium/grid/distributor/httpd/DistributorServer.java

@@ -106,7 +106,8 @@ public Executable configure(String... args) {
       Distributor distributor = new LocalDistributor(
           tracer,
           bus,
-          HttpClient.Factory.createDefault());
+          HttpClient.Factory.createDefault(),
+          null);
 
       BaseServerOptions serverOptions = new BaseServerOptions(config);
 

java/server/src/org/openqa/selenium/grid/distributor/local/BUCK

@@ -12,6 +12,7 @@ java_library(
         "//java/server/src/org/openqa/selenium/grid/web:web",
         "//java/server/src/org/openqa/selenium/grid/node:node",
         "//java/server/src/org/openqa/selenium/grid/node/remote:remote",
+        "//java/server/src/org/openqa/selenium/grid/sessionmap:sessionmap",
         "//third_party/java/guava:guava",
     ],
     visibility = [

java/server/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java

@@ -33,6 +33,7 @@
 import org.openqa.selenium.grid.distributor.Distributor;
 import org.openqa.selenium.grid.distributor.DistributorStatus;
 import org.openqa.selenium.grid.node.Node;
+import org.openqa.selenium.grid.sessionmap.SessionMap;
 import org.openqa.selenium.json.Json;
 import org.openqa.selenium.json.JsonOutput;
 import org.openqa.selenium.remote.NewSessionPayload;
@@ -66,16 +67,19 @@ public class LocalDistributor extends Distributor {
   private final Set<Host> hosts = new HashSet<>();
   private final DistributedTracer tracer;
   private final EventBus bus;
+  private final SessionMap sessions;
   private final Regularly hostChecker = new Regularly("distributor host checker");
   private final Map<UUID, Collection<Runnable>> allChecks = new ConcurrentHashMap<>();
 
   public LocalDistributor(
       DistributedTracer tracer,
       EventBus bus,
-      HttpClient.Factory httpClientFactory) {
+      HttpClient.Factory httpClientFactory,
+      SessionMap sessions) {
     super(tracer, httpClientFactory);
     this.tracer = Objects.requireNonNull(tracer);
     this.bus = Objects.requireNonNull(bus);
+    this.sessions = Objects.requireNonNull(sessions);
   }
 
   @Override
@@ -107,10 +111,14 @@ public Session newSession(NewSessionPayload payload) throws SessionNotCreatedExc
       writeLock.unlock();
     }
 
-    return selected
+    Session session = selected
         .orElseThrow(
             () -> new SessionNotCreatedException("Unable to find provider for session: " + allCaps))
         .get();
+
+    sessions.add(session);
+
+    return session;
   }
 
   @Override

java/server/src/org/openqa/selenium/grid/node/httpd/NodeServer.java

@@ -131,8 +131,7 @@ public Executable configure(String... args) {
           tracer,
           bus,
           httpClientFactory,
-          serverOptions.getExternalUri(),
-          sessions);
+          serverOptions.getExternalUri());
       nodeFlags.configure(config, httpClientFactory, builder);
       LocalNode node = builder.build();
 

java/server/src/org/openqa/selenium/grid/node/local/BUCK

@@ -8,7 +8,6 @@ java_library(
         "//java/server/src/org/openqa/selenium/events:events",
         "//java/server/src/org/openqa/selenium/grid/config:config",
         "//java/server/src/org/openqa/selenium/grid/node:node",
-        "//java/server/src/org/openqa/selenium/grid/sessionmap:sessionmap",
         "//java/server/src/org/openqa/selenium/grid/web:web",
         "//third_party/java/beust:jcommander",
         "//third_party/java/guava:guava",

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

@@ -36,7 +36,7 @@
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.grid.data.SessionClosedEvent;
 import org.openqa.selenium.grid.node.Node;
-import org.openqa.selenium.grid.sessionmap.SessionMap;
+import org.openqa.selenium.grid.web.CommandHandler;
 import org.openqa.selenium.remote.SessionId;
 import org.openqa.selenium.remote.http.HttpClient;
 import org.openqa.selenium.remote.http.HttpRequest;
@@ -154,7 +154,7 @@ public Optional<Session> newSession(Capabilities capabilities) {
 
       // The session we return has to look like it came from the node, since we might be dealing
       // with a webdriver implementation that only accepts connections from localhost
-      return Optional.of(new Session(session.getId(), externalUri, session.getCapabilities()));
+      return Optional.of(createExternalSession(session, externalUri));
     }
   }
 
@@ -183,7 +183,7 @@ public Session getSession(SessionId id) throws NoSuchSessionException {
 
       span.addTag("session.capabilities", session.getCapabilities());
       span.addTag("session.uri", session.getUri());
-      return new Session(session.getId(), externalUri, session.getCapabilities());
+      return createExternalSession(session, externalUri);
     }
   }
 
@@ -237,13 +237,20 @@ public void stop(SessionId id) throws NoSuchSessionException {
     }
   }
 
+  private Session createExternalSession(SessionAndHandler other, URI externalUri) {
+    return new HandledSession(
+      other.getId(),
+      externalUri,
+      other.getCapabilities(),
+      other.getHandler());
+  }
+
   private void killSession(Span span, SessionAndHandler session) {
     span.addTag("session.id", session.getId());
     span.addTag("session.capabilities", session.getCapabilities());
     span.addTag("session.uri", session.getUri());
 
     currentSessions.invalidate(session.getId());
-    session.stop();
     // Attempt to stop the session
     session.stop();
     bus.fire(new SessionClosedEvent(session.getId()));
@@ -289,9 +296,8 @@ public static Builder builder(
       DistributedTracer tracer,
       EventBus bus,
       HttpClient.Factory httpClientFactory,
-      URI uri,
-      SessionMap sessions) {
-    return new Builder(tracer, bus, httpClientFactory, uri, sessions);
+      URI uri) {
+    return new Builder(tracer, bus, httpClientFactory, uri);
   }
 
   public static class Builder {
@@ -300,7 +306,6 @@ public static class Builder {
     private final EventBus bus;
     private final HttpClient.Factory httpClientFactory;
     private final URI uri;
-    private final SessionMap sessions;
     private final ImmutableList.Builder<SessionFactory> factories;
     private int maxCount = Runtime.getRuntime().availableProcessors() * 5;
     private Ticker ticker = Ticker.systemTicker();
@@ -311,21 +316,19 @@ public Builder(
         DistributedTracer tracer,
         EventBus bus,
         HttpClient.Factory httpClientFactory,
-        URI uri,
-        SessionMap sessions) {
+        URI uri) {
       this.tracer = Objects.requireNonNull(tracer);
       this.bus = Objects.requireNonNull(bus);
       this.httpClientFactory = Objects.requireNonNull(httpClientFactory);
       this.uri = Objects.requireNonNull(uri);
-      this.sessions = Objects.requireNonNull(sessions);
       this.factories = ImmutableList.builder();
     }
 
     public Builder add(Capabilities stereotype, Function<Capabilities, Session> factory) {
       Objects.requireNonNull(stereotype, "Capabilities must be set.");
       Objects.requireNonNull(factory, "Session factory must be set.");
 
-      factories.add(new SessionFactory(httpClientFactory, sessions, stereotype, factory));
+      factories.add(new SessionFactory(httpClientFactory, stereotype, factory));
 
       return this;
     }
@@ -388,4 +391,18 @@ public Node build() {
     }
   }
 
+  class HandledSession extends Session implements CommandHandler {
+
+    private final CommandHandler handler;
+
+    public HandledSession(SessionId id, URI uri, Capabilities caps, CommandHandler handler) {
+      super(id, uri, caps);
+      this.handler = Objects.requireNonNull(handler);
+    }
+
+    @Override
+    public void execute(HttpRequest req, HttpResponse resp) throws IOException {
+      handler.execute(req, resp);
+    }
+  }
 }

java/server/src/org/openqa/selenium/grid/node/local/SessionFactory.java

@@ -22,7 +22,6 @@
 import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.ImmutableCapabilities;
 import org.openqa.selenium.grid.data.Session;
-import org.openqa.selenium.grid.sessionmap.SessionMap;
 import org.openqa.selenium.grid.web.CommandHandler;
 import org.openqa.selenium.grid.web.ReverseProxyHandler;
 import org.openqa.selenium.remote.http.HttpClient;
@@ -38,18 +37,15 @@ class SessionFactory
     implements Predicate<Capabilities>, Function<Capabilities, Optional<SessionAndHandler>> {
 
   private final HttpClient.Factory httpClientFactory;
-  private final SessionMap sessions;
   private final Capabilities capabilities;
   private final Function<Capabilities, Session> generator;
   private volatile boolean available = true;
 
   SessionFactory(
       HttpClient.Factory httpClientFactory,
-      SessionMap sessions,
       Capabilities capabilities,
       Function<Capabilities, Session> generator) {
     this.httpClientFactory = Objects.requireNonNull(httpClientFactory);
-    this.sessions = Objects.requireNonNull(sessions);
     this.capabilities = Objects.requireNonNull(ImmutableCapabilities.copyOf(capabilities));
     this.generator = Objects.requireNonNull(generator);
   }
@@ -88,7 +84,6 @@ public Optional<SessionAndHandler> apply(Capabilities capabilities) {
       this.available = true;
       return Optional.empty();
     }
-    sessions.add(session);
 
     CommandHandler handler;
     if (session instanceof CommandHandler) {
@@ -104,14 +99,10 @@ public Optional<SessionAndHandler> apply(Capabilities capabilities) {
 
     String killUrl = "/session/" + session.getId();
     CommandHandler killingHandler = (req, res) -> {
+      handler.execute(req, res);
       if (req.getMethod() == DELETE && killUrl.equals(req.getUri())) {
-        try {
-          sessions.remove(session.getId());
-        } finally {
-          available = true;
-        }
+        available = true;
       }
-      handler.execute(req, res);
     };
 
     return Optional.of(new SessionAndHandler(session, killingHandler));