[grid] remove indirection in cache structures

Author: joerg1985

java/src/org/openqa/selenium/grid/node/Node.java

@@ -32,7 +32,6 @@
 import java.util.Map;
 import java.util.ServiceLoader;
 import java.util.Set;
-import java.util.UUID;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
@@ -245,7 +244,7 @@ public TemporaryFilesystem getUploadsFilesystem(SessionId id) throws IOException
     throw new UnsupportedOperationException();
   }
 
-  public TemporaryFilesystem getDownloadsFilesystem(UUID uuid) throws IOException {
+  public TemporaryFilesystem getDownloadsFilesystem(SessionId id) throws IOException {
     throw new UnsupportedOperationException();
   }
 

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

@@ -139,8 +139,7 @@ public class LocalNode extends Node implements Closeable {
   private final List<SessionSlot> factories;
   private final Cache<SessionId, SessionSlot> currentSessions;
   private final Cache<SessionId, TemporaryFilesystem> uploadsTempFileSystem;
-  private final Cache<UUID, TemporaryFilesystem> downloadsTempFileSystem;
-  private final Cache<SessionId, UUID> sessionToDownloadsDir;
+  private final Cache<SessionId, TemporaryFilesystem> downloadsTempFileSystem;
   private final AtomicInteger pendingSessions = new AtomicInteger();
   private final AtomicInteger sessionCount = new AtomicInteger();
   private final Runnable shutdown;
@@ -198,7 +197,7 @@ protected LocalNode(
             : healthCheck;
 
     // Do not clear this cache automatically using a timer.
-    // It will be explicitly cleaned up, as and when "sessionToDownloadsDir" is auto cleaned.
+    // It will be explicitly cleaned up, as and when "currentSessions" is auto cleaned.
     this.uploadsTempFileSystem =
         CacheBuilder.newBuilder()
             .removalListener(
@@ -213,11 +212,11 @@ protected LocalNode(
             .build();
 
     // Do not clear this cache automatically using a timer.
-    // It will be explicitly cleaned up, as and when "sessionToDownloadsDir" is auto cleaned.
+    // It will be explicitly cleaned up, as and when "currentSessions" is auto cleaned.
     this.downloadsTempFileSystem =
         CacheBuilder.newBuilder()
             .removalListener(
-                (RemovalListener<UUID, TemporaryFilesystem>)
+                (RemovalListener<SessionId, TemporaryFilesystem>)
                     notification ->
                         Optional.ofNullable(notification.getValue())
                             .ifPresent(
@@ -227,32 +226,6 @@ protected LocalNode(
                                 }))
             .build();
 
-    // Do not clear this cache automatically using a timer.
-    // It will be explicitly cleaned up, as and when "currentSessions" is auto cleaned.
-    this.sessionToDownloadsDir =
-        CacheBuilder.newBuilder()
-            .removalListener(
-                (RemovalListener<SessionId, UUID>)
-                    notification -> {
-                      Optional.ofNullable(notification.getValue())
-                          .ifPresent(
-                              value -> {
-                                downloadsTempFileSystem.invalidate(value);
-                                LOG.fine(
-                                    "Removing Downloads folder associated with "
-                                        + notification.getKey());
-                              });
-                      Optional.ofNullable(notification.getKey())
-                          .ifPresent(
-                              value -> {
-                                uploadsTempFileSystem.invalidate(value);
-                                LOG.fine(
-                                    "Removing Uploads folder associated with "
-                                        + notification.getKey());
-                              });
-                    })
-            .build();
-
     this.currentSessions =
         CacheBuilder.newBuilder()
             .expireAfterAccess(sessionTimeout)
@@ -386,15 +359,6 @@ public int getCurrentSessionCount() {
     return Math.toIntExact(n);
   }
 
-  @VisibleForTesting
-  public UUID getDownloadsIdForSession(SessionId id) {
-    UUID uuid = sessionToDownloadsDir.getIfPresent(id);
-    if (uuid == null) {
-      throw new NoSuchSessionException("Cannot find session with id: " + id);
-    }
-    return uuid;
-  }
-
   @ManagedAttribute(name = "MaxSessions")
   public int getMaxSessionCount() {
     return maxSessionCount;
@@ -508,21 +472,32 @@ public Either<WebDriverException, CreateSessionResponse> newSession(
         return Either.left(new RetrySessionRequestException("Drain after session count reached."));
       }
 
-      UUID uuidForSessionDownloads = UUID.randomUUID();
       Capabilities desiredCapabilities = sessionRequest.getDesiredCapabilities();
+      TemporaryFilesystem downloadsTfs;
       if (managedDownloadsRequested(desiredCapabilities)) {
-        Capabilities enhanced = setDownloadsDirectory(uuidForSessionDownloads, desiredCapabilities);
+        UUID uuidForSessionDownloads = UUID.randomUUID();
+
+        downloadsTfs =
+            TemporaryFilesystem.getTmpFsBasedOn(
+                TemporaryFilesystem.getDefaultTmpFS()
+                    .createTempDir("uuid", uuidForSessionDownloads.toString()));
+
+        Capabilities enhanced = setDownloadsDirectory(downloadsTfs, desiredCapabilities);
         enhanced = desiredCapabilities.merge(enhanced);
         sessionRequest =
             new CreateSessionRequest(
                 sessionRequest.getDownstreamDialects(), enhanced, sessionRequest.getMetadata());
+      } else {
+        downloadsTfs = null;
       }
 
       Either<WebDriverException, ActiveSession> possibleSession = slotToUse.apply(sessionRequest);
 
       if (possibleSession.isRight()) {
         ActiveSession session = possibleSession.right();
-        sessionToDownloadsDir.put(session.getId(), uuidForSessionDownloads);
+        if (downloadsTfs != null) {
+          downloadsTempFileSystem.put(session.getId(), downloadsTfs);
+        }
         currentSessions.put(session.getId(), slotToUse);
 
         SessionId sessionId = session.getId();
@@ -558,6 +533,10 @@ public Either<WebDriverException, CreateSessionResponse> newSession(
                 getEncoder(session.getDownstreamDialect()).apply(externalSession)));
       } else {
         slotToUse.release();
+        if (downloadsTfs != null) {
+          downloadsTfs.deleteTemporaryFiles();
+          downloadsTfs.deleteBaseDir();
+        }
         span.setAttribute(AttributeKey.ERROR.getKey(), true);
         span.setStatus(Status.ABORTED);
         span.addEvent("Unable to create session with the driver", attributeMap);
@@ -576,15 +555,8 @@ private boolean managedDownloadsRequested(Capabilities capabilities) {
         && Boolean.parseBoolean(downloadsEnabled.toString());
   }
 
-  private Capabilities setDownloadsDirectory(UUID uuid, Capabilities caps) {
-    File tempDir;
-    try {
-      TemporaryFilesystem tempFS = getDownloadsFilesystem(uuid);
-      //      tempDir = tempFS.createTempDir("download", "file");
-      tempDir = tempFS.createTempDir("download", "");
-    } catch (IOException e) {
-      throw new UncheckedIOException(e);
-    }
+  private Capabilities setDownloadsDirectory(TemporaryFilesystem downloadsTfs, Capabilities caps) {
+    File tempDir = downloadsTfs.createTempDir("download", "");
     if (Browser.CHROME.is(caps) || Browser.EDGE.is(caps)) {
       ImmutableMap<String, Serializable> map =
           ImmutableMap.of(
@@ -700,16 +672,8 @@ public TemporaryFilesystem getUploadsFilesystem(SessionId id) throws IOException
   }
 
   @Override
-  public TemporaryFilesystem getDownloadsFilesystem(UUID uuid) throws IOException {
-    try {
-      return downloadsTempFileSystem.get(
-          uuid,
-          () ->
-              TemporaryFilesystem.getTmpFsBasedOn(
-                  TemporaryFilesystem.getDefaultTmpFS().createTempDir("uuid", uuid.toString())));
-    } catch (ExecutionException e) {
-      throw new IOException(e);
-    }
+  public TemporaryFilesystem getDownloadsFilesystem(SessionId sessionId) throws IOException {
+    return downloadsTempFileSystem.getIfPresent(sessionId);
   }
 
   @Override
@@ -746,11 +710,7 @@ public HttpResponse downloadFile(HttpRequest req, SessionId id) {
               + "[--enable-managed-downloads] and restart the node";
       throw new WebDriverException(msg);
     }
-    UUID uuid = sessionToDownloadsDir.getIfPresent(id);
-    if (uuid == null) {
-      throw new NoSuchSessionException("Cannot find session with id: " + id);
-    }
-    TemporaryFilesystem tempFS = downloadsTempFileSystem.getIfPresent(uuid);
+    TemporaryFilesystem tempFS = downloadsTempFileSystem.getIfPresent(id);
     if (tempFS == null) {
       String msg =
           "Cannot find downloads file system for session id: "
@@ -860,8 +820,11 @@ public HttpResponse uploadFile(HttpRequest req, SessionId id) {
   public void stop(SessionId id) throws NoSuchSessionException {
     Require.nonNull("Session ID", id);
 
-    if (sessionToDownloadsDir.getIfPresent(id) != null) {
-      sessionToDownloadsDir.invalidate(id);
+    if (downloadsTempFileSystem.getIfPresent(id) != null) {
+      downloadsTempFileSystem.invalidate(id);
+    }
+    if (uploadsTempFileSystem.getIfPresent(id) != null) {
+      uploadsTempFileSystem.invalidate(id);
     }
 
     SessionSlot slot = currentSessions.getIfPresent(id);

java/test/org/openqa/selenium/grid/node/NodeTest.java

@@ -640,8 +640,8 @@ void canDownloadMultipleFile() throws IOException {
       String decodedContents = String.join("", Files.readAllLines(path));
       assertThat(decodedContents).isEqualTo(hello);
     } finally {
-      UUID downloadsId = local.getDownloadsIdForSession(session.getId());
-      File someDir = getTemporaryFilesystemBaseDir(local.getDownloadsFilesystem(downloadsId));
+      TemporaryFilesystem downloadsTfs = local.getDownloadsFilesystem(session.getId());
+      File someDir = getTemporaryFilesystemBaseDir(downloadsTfs);
       node.stop(session.getId());
       assertThat(someDir).doesNotExist();
     }
@@ -947,8 +947,8 @@ private CreateSessionRequest createSessionRequest(Capabilities caps) {
 
   private String simulateFileDownload(SessionId id, String text) throws IOException {
     File zip = createTmpFile(text);
-    UUID downloadsId = local.getDownloadsIdForSession(id);
-    File someDir = getTemporaryFilesystemBaseDir(local.getDownloadsFilesystem(downloadsId));
+    TemporaryFilesystem downloadsTfs = local.getDownloadsFilesystem(id);
+    File someDir = getTemporaryFilesystemBaseDir(downloadsTfs);
     File downloadsDirectory = Optional.ofNullable(someDir.listFiles()).orElse(new File[] {})[0];
     File target = new File(downloadsDirectory, zip.getName());
     boolean renamed = zip.renameTo(target);