Fix JWT refresh and use TowerClient to handle Fusion token requests(#6207)

Signed-off-by: Jordi Deu-Pons <[email protected]>
Signed-off-by: Paolo Di Tommaso <[email protected]>
Co-authored-by: Paolo Di Tommaso <[email protected]>

Author: jordeu

plugins/nf-tower/src/main/io/seqera/tower/plugin/TowerFactory.groovy

@@ -100,13 +100,13 @@ class TowerFactory implements TraceObserverFactory {
     @Memoized
     static TowerClient client(Session session, Map<String,String> env) {
         final config = session.config
-        Boolean isEnabled = config.navigate('tower.enabled') as Boolean || env.get('TOWER_WORKFLOW_ID')
+        Boolean isEnabled = config.navigate('tower.enabled') as Boolean || env.get('TOWER_WORKFLOW_ID') || config.navigate('fusion.enabled')
         return isEnabled
             ? createTowerClient0(session, config, env)
             : null
     }
 
     static TowerClient client() {
-        client(Global.session as Session, SysEnv.get())
+        return client(Global.session as Session, SysEnv.get())
     }
 }

plugins/nf-tower/src/main/io/seqera/tower/plugin/TowerFusionToken.groovy

@@ -1,31 +1,15 @@
 package io.seqera.tower.plugin
 
-import java.net.http.HttpClient
-import java.net.http.HttpRequest
-import java.net.http.HttpResponse
-import java.time.Duration
-import java.time.Instant
-import java.time.temporal.ChronoUnit
-import java.util.concurrent.Executors
-import java.util.function.Predicate
-
 import com.google.common.cache.Cache
 import com.google.common.cache.CacheBuilder
 import com.google.common.util.concurrent.UncheckedExecutionException
-import com.google.gson.Gson
 import com.google.gson.JsonSyntaxException
-import dev.failsafe.Failsafe
-import dev.failsafe.RetryPolicy
-import dev.failsafe.event.EventListener
-import dev.failsafe.event.ExecutionAttemptedEvent
-import dev.failsafe.function.CheckedSupplier
 import groovy.transform.CompileStatic
 import groovy.util.logging.Slf4j
 import io.seqera.tower.plugin.exception.BadResponseException
 import io.seqera.tower.plugin.exception.UnauthorizedException
 import io.seqera.tower.plugin.exchange.GetLicenseTokenRequest
 import io.seqera.tower.plugin.exchange.GetLicenseTokenResponse
-import io.seqera.util.trace.TraceUtils
 import nextflow.SysEnv
 import nextflow.exception.AbortOperationException
 import nextflow.exception.ReportWarningException
@@ -34,8 +18,12 @@ import nextflow.fusion.FusionToken
 import nextflow.platform.PlatformHelper
 import nextflow.plugin.Priority
 import nextflow.util.GsonHelper
-import nextflow.util.Threads
 import org.pf4j.Extension
+
+import java.time.Duration
+import java.time.Instant
+import java.time.temporal.ChronoUnit
+
 /**
  * Environment provider for Platform-specific environment variables.
  *
@@ -50,24 +38,6 @@ class TowerFusionToken implements FusionToken {
     // The path relative to the Platform endpoint where license-scoped JWT tokens are obtained
     private static final String LICENSE_TOKEN_PATH = 'license/token/'
 
-    // Server errors that should trigger a retry
-    private static final List<Integer> SERVER_ERRORS = [408, 429, 500, 502, 503, 504]
-
-    // Default connection timeout for HTTP requests
-    private static final Duration DEFAULT_CONNECTION_TIMEOUT = Duration.of(30, ChronoUnit.SECONDS)
-
-    // Default retry policy settings for HTTP requests: delay, max delay, attempts, and jitter
-    private static final Duration DEFAULT_RETRY_POLICY_DELAY = Duration.of(450, ChronoUnit.MILLIS)
-    private static final Duration DEFAULT_RETRY_POLICY_MAX_DELAY = Duration.of(90, ChronoUnit.SECONDS)
-    private static final int DEFAULT_RETRY_POLICY_MAX_ATTEMPTS = 10
-    private static final double DEFAULT_RETRY_POLICY_JITTER = 0.5
-
-    // The HttpClient instance used to send requests
-    private final HttpClient httpClient = newDefaultHttpClient()
-
-    // The RetryPolicy instance used to retry requests
-    private final RetryPolicy retryPolicy = newDefaultRetryPolicy(SERVER_ERRORS)
-
     // Time-to-live for cached tokens
     private Duration tokenTTL = Duration.of(1, ChronoUnit.HOURS)
 
@@ -76,32 +46,21 @@ class TowerFusionToken implements FusionToken {
         .expireAfterWrite(tokenTTL)
         .build()
 
-    // Platform endpoint to use for requests
-    private String endpoint
-
-    // Platform access token to use for requests
-    private String accessToken
-
     // Platform workflowId
     private String workspaceId
 
     // Platform workflowId
     private String workflowId
 
+    // Platform client to handle all the requests
+    private TowerClient client
+
     TowerFusionToken() {
         final config = PlatformHelper.config()
         final env = SysEnv.get()
-        this.endpoint = PlatformHelper.getEndpoint(config, env)
-        this.accessToken = PlatformHelper.getAccessToken(config, env)
         this.workflowId = env.get('TOWER_WORKFLOW_ID')
         this.workspaceId = PlatformHelper.getWorkspaceId(config, env)
-    }
-
-    protected void validateConfig() {
-        if( !endpoint )
-            throw new IllegalArgumentException("Missing Seqera Platform endpoint")
-        if( !accessToken )
-            throw new IllegalArgumentException("Missing Seqera Platform access token")
+        this.client = TowerFactory.client()
     }
 
     /**
@@ -125,7 +84,6 @@ class TowerFusionToken implements FusionToken {
     }
 
     protected Map<String,String> getEnvironment0(String scheme, FusionConfig config) {
-        validateConfig()
         final product = config.sku()
         final version = config.version()
         final token = getLicenseToken(product, version)
@@ -178,102 +136,6 @@ class TowerFusionToken implements FusionToken {
      * Helper methods
      *************************************************************************/
 
-    /**
-     * Create a new HttpClient instance with default settings
-     * @return The new HttpClient instance
-     */
-    private static HttpClient newDefaultHttpClient() {
-        final builder = HttpClient.newBuilder()
-            .version(HttpClient.Version.HTTP_1_1)
-            .followRedirects(HttpClient.Redirect.NEVER)
-            .cookieHandler(new CookieManager())
-            .connectTimeout(DEFAULT_CONNECTION_TIMEOUT)
-        // use virtual threads executor if enabled
-        if ( Threads.useVirtual() ) {
-            builder.executor(Executors.newVirtualThreadPerTaskExecutor())
-        }
-        // build and return the new client
-        return builder.build()
-    }
-
-    /**
-     * Create a new RetryPolicy instance with default settings and the given list of retryable errors. With this policy,
-     * a request is retried on IOExceptions and any server errors defined in errorsToRetry. The number of retries, delay,
-     * max delay, and jitter are controlled by the corresponding values defined at class level.
-     *
-     * @return The new RetryPolicy instance
-     */
-    private static <T> RetryPolicy<HttpResponse<T>> newDefaultRetryPolicy(List<Integer> errorsToRetry) {
-
-        final retryOnException = (e -> e instanceof IOException) as Predicate<? extends Throwable>
-        final retryOnStatusCode = ((HttpResponse<T> resp) -> resp.statusCode() in errorsToRetry) as Predicate<HttpResponse<T>>
-
-        final listener = new EventListener<ExecutionAttemptedEvent<HttpResponse<T>>>() {
-            @Override
-            void accept(ExecutionAttemptedEvent event) throws Throwable {
-                def msg = "connection failure - attempt: ${event.attemptCount}"
-                if (event.lastResult != null)
-                    msg += "; response: ${event.lastResult}"
-                if (event.lastFailure != null)
-                    msg += "; exception: [${event.lastFailure.class.name}] ${event.lastFailure.message}"
-                log.debug(msg)
-            }
-        }
-        return RetryPolicy.<HttpResponse<T>> builder()
-            .handleIf(retryOnException)
-            .handleResultIf(retryOnStatusCode)
-            .withBackoff(DEFAULT_RETRY_POLICY_DELAY.toMillis(), DEFAULT_RETRY_POLICY_MAX_DELAY.toMillis(), ChronoUnit.MILLIS)
-            .withMaxAttempts(DEFAULT_RETRY_POLICY_MAX_ATTEMPTS)
-            .withJitter(DEFAULT_RETRY_POLICY_JITTER)
-            .onRetry(listener)
-            .build()
-    }
-
-    /**
-     * Send an HTTP request and return the response. This method automatically retries the request according to the
-     * given RetryPolicy.
-     *
-     * @param req The HttpRequest to send
-     * @return The HttpResponse received
-     */
-    private <T> HttpResponse<String> safeHttpSend(HttpRequest req, RetryPolicy<T> policy) {
-        return Failsafe.with(policy).get(
-            () -> {
-                log.debug "Http request: method=${req.method()}; uri=${req.uri()}; request=${req}"
-                final resp = httpClient.send(req, HttpResponse.BodyHandlers.ofString())
-                log.debug "Http response: statusCode=${resp.statusCode()}; body=${resp.body()}"
-                return resp
-            } as CheckedSupplier
-        ) as HttpResponse<String>
-    }
-
-    /**
-     * Create a {@link HttpRequest} representing a {@link GetLicenseTokenRequest} object
-     *
-     * @param req The LicenseTokenRequest object
-     * @return The resulting HttpRequest object
-     */
-    private HttpRequest makeHttpRequest(GetLicenseTokenRequest req) {
-        final body = HttpRequest.BodyPublishers.ofString( GsonHelper.toJson(req) )
-        return HttpRequest.newBuilder()
-            .uri(URI.create("${endpoint}/${LICENSE_TOKEN_PATH}").normalize())
-            .header('Content-Type', 'application/json')
-            .header('Traceparent', TraceUtils.rndTrace())
-            .header('Authorization', "Bearer ${accessToken}")
-            .POST(body)
-            .build()
-    }
-
-    /**
-     * Serialize a {@link GetLicenseTokenRequest} object into a JSON string
-     *
-     * @param req The LicenseTokenRequest object
-     * @return The resulting JSON string
-     */
-    private static String serializeToJson(GetLicenseTokenRequest req) {
-        return new Gson().toJson(req)
-    }
-
     /**
      * Parse a JSON string into a {@link GetLicenseTokenResponse} object
      *
@@ -299,24 +161,18 @@ class TowerFusionToken implements FusionToken {
      */
     private GetLicenseTokenResponse sendRequest(GetLicenseTokenRequest req) throws AbortOperationException, UnauthorizedException, BadResponseException, IllegalStateException {
 
-        final httpReq = makeHttpRequest(req)
-
-        try {
-            final resp = safeHttpSend(httpReq, retryPolicy)
-
-            if( resp.statusCode() == 200 ) {
-                final ret = parseLicenseTokenResponse(resp.body())
-                return ret
-            }
+        final url = "${client.getEndpoint()}/${LICENSE_TOKEN_PATH}"
+        final resp = client.sendHttpMessage(url, req.toMap())
 
-            if( resp.statusCode() == 401 ) {
-                throw new UnauthorizedException("Unauthorized [401] - Verify you have provided a Seqera Platform valid access token")
-            }
-
-            throw new BadResponseException("Invalid response: ${httpReq.method()} ${httpReq.uri()} [${resp.statusCode()}] ${resp.body()}")
+        if( resp.code == 200 ) {
+            final ret = parseLicenseTokenResponse(resp.message)
+            return ret
         }
-        catch (IOException e) {
-            throw new IllegalStateException("Unable to send request to '${httpReq.uri()}' : ${e.message}")
+
+        if( resp.code == 401 ) {
+            throw new UnauthorizedException("Unauthorized [401] - Verify you have provided a Seqera Platform valid access token")
         }
+
+        throw new BadResponseException("Invalid response: ${url} [${resp.code}] ${resp.message} -- ${resp.cause}")
     }
 }

plugins/nf-tower/src/main/io/seqera/tower/plugin/exchange/GetLicenseTokenRequest.groovy

@@ -29,4 +29,16 @@ class GetLicenseTokenRequest {
      * The Platform workspace ID associated with this request
      */
     String workspaceId
+
+    /**
+     * @return a Map representation of the request
+     */
+    Map<String, String> toMap() {
+        final map = new HashMap<String, String>()
+        map.product = this.product
+        map.version = this.version
+        map.workflowId = this.workflowId
+        map.workspaceId = this.workspaceId
+        return map
+    }
 }