Prefetch all jdks in gradle.py

Change-Id: I608018def8eddba6cbf68d121442160858994868
diff --git a/tools/gradle.py b/tools/gradle.py
index 8a80347..00aca65 100755
--- a/tools/gradle.py
+++ b/tools/gradle.py
@@ -68,10 +68,12 @@
     get_gradle(True), GRADLE8_TGZ, GRADLE8_SHA1, 'Gradle binary')
 
 def EnsureJdk():
-  jdkRoot = jdk.GetJdkRoot()
-  jdkTgz = jdkRoot + '.tar.gz'
-  jdkSha1 = jdkTgz + '.sha1'
-  utils.EnsureDepFromGoogleCloudStorage(jdkRoot, jdkTgz, jdkSha1, 'JDK')
+  # Gradle in the new setup will use the jdks in the evaluation - fetch
+  # all beforehand.
+  for root in jdk.GetAllJdkDirs():
+    jdkTgz = root + '.tar.gz'
+    jdkSha1 = jdkTgz + '.sha1'
+    utils.EnsureDepFromGoogleCloudStorage(root, jdkTgz, jdkSha1, root)
 
 def EnsureDeps():
   EnsureGradle()
diff --git a/tools/jdk.py b/tools/jdk.py
index 5cf1465..f37d272 100755
--- a/tools/jdk.py
+++ b/tools/jdk.py
@@ -10,6 +10,10 @@
 
 JDK_DIR = os.path.join(defines.THIRD_PARTY, 'openjdk')
 
+ALL_JDKS = ['openjdk-9.0.4', 'jdk-11', 'jdk-15', 'jdk-16', 'jdk-17',
+            'jdk-18', 'jdk-20']
+
+
 def GetJdkHome():
   return GetJdk11Home()
 
@@ -18,6 +22,10 @@
 
 def GetJdk11Root():
   root = os.path.join(JDK_DIR, 'jdk-11')
+  os_root = GetOSPath(root)
+  return os_root if os_root else os.environ['JAVA_HOME']
+
+def GetOSPath(root):
   if defines.IsLinux():
     return os.path.join(root, 'linux')
   elif defines.IsOsX():
@@ -25,7 +33,16 @@
   elif defines.IsWindows():
     return os.path.join(root, 'windows')
   else:
-    return os.environ['JAVA_HOME']
+    return None
+
+def GetAllJdkDirs():
+  dirs = []
+  for jdk in ALL_JDKS:
+    root = GetOSPath(os.path.join(JDK_DIR, jdk))
+    # Some jdks are not available on windows, don't try to get these.
+    if os.path.exists(root + '.tar.gz.sha1'):
+      dirs.append(root)
+  return dirs
 
 def GetJdk11Home():
   root = GetJdk11Root()