[grid] Improve readTimeout in handle session between Router and Node

Signed-off-by: Viet Nguyen Duc <[email protected]>

Author: VietND96

java/src/org/openqa/selenium/grid/router/HandleSession.java

@@ -21,13 +21,17 @@
 import static org.openqa.selenium.remote.RemoteTags.SESSION_ID;
 import static org.openqa.selenium.remote.RemoteTags.SESSION_ID_EVENT;
 import static org.openqa.selenium.remote.http.Contents.asJson;
+import static org.openqa.selenium.remote.http.Contents.reader;
+import static org.openqa.selenium.remote.http.HttpMethod.GET;
 import static org.openqa.selenium.remote.tracing.Tags.EXCEPTION;
 import static org.openqa.selenium.remote.tracing.Tags.HTTP_REQUEST;
 import static org.openqa.selenium.remote.tracing.Tags.HTTP_REQUEST_EVENT;
 import static org.openqa.selenium.remote.tracing.Tags.HTTP_RESPONSE;
 
 import java.io.Closeable;
+import java.io.Reader;
 import java.net.URI;
+import java.time.Duration;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.Iterator;
@@ -43,9 +47,12 @@
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.concurrent.ExecutorServices;
 import org.openqa.selenium.concurrent.GuardedRunnable;
+import org.openqa.selenium.grid.data.NodeStatus;
 import org.openqa.selenium.grid.sessionmap.SessionMap;
 import org.openqa.selenium.grid.web.ReverseProxyHandler;
 import org.openqa.selenium.internal.Require;
+import org.openqa.selenium.json.Json;
+import org.openqa.selenium.json.JsonInput;
 import org.openqa.selenium.remote.ErrorCodec;
 import org.openqa.selenium.remote.SessionId;
 import org.openqa.selenium.remote.http.ClientConfig;
@@ -229,8 +236,7 @@ private Callable<UsageCountingReverseProxyHandler> loadSessionId(
                       return entry;
                     }
 
-                    ClientConfig config =
-                        ClientConfig.defaultConfig().baseUri(sessionUri).withRetries();
+                    ClientConfig config = fetchNodeSessionTimeout(sessionUri);
                     HttpClient httpClient = httpClientFactory.createClient(config);
 
                     return new CacheEntry(httpClient, 1);
@@ -246,6 +252,49 @@ private Callable<UsageCountingReverseProxyHandler> loadSessionId(
         });
   }
 
+  private ClientConfig fetchNodeSessionTimeout(URI uri) {
+    ClientConfig config = ClientConfig.defaultConfig().baseUri(uri).withRetries();
+    Duration sessionTimeout = config.readTimeout();
+    try (HttpClient httpClient = httpClientFactory.createClient(config)) {
+      HttpRequest statusRequest = new HttpRequest(GET, "/status");
+      HttpResponse res = httpClient.execute(statusRequest);
+      Reader reader = reader(res);
+      Json JSON = new Json();
+      JsonInput in = JSON.newInput(reader);
+      in.beginObject();
+      // Skip everything until we find "value"
+      while (in.hasNext()) {
+        if ("value".equals(in.nextName())) {
+          in.beginObject();
+          while (in.hasNext()) {
+            if ("node".equals(in.nextName())) {
+              NodeStatus nodeStatus = in.read(NodeStatus.class);
+              sessionTimeout = nodeStatus.getSessionTimeout();
+              LOG.fine(
+                  "Fetched session timeout from node status (read timeout: "
+                      + sessionTimeout.toSeconds()
+                      + " seconds) for "
+                      + uri);
+            } else {
+              in.skipValue();
+            }
+          }
+          in.endObject();
+        } else {
+          in.skipValue();
+        }
+      }
+    } catch (Exception e) {
+      LOG.fine(
+          "Use default from ClientConfig (read timeout: "
+              + config.readTimeout().toSeconds()
+              + " seconds) for "
+              + uri);
+    }
+    config = config.readTimeout(sessionTimeout);
+    return config;
+  }
+
   @Override
   public void close() {
     ExecutorServices.shutdownGracefully(