binarycat0 commented on code in PR #90: URL: https://github.com/apache/polaris-tools/pull/90#discussion_r2606872951
########## benchmarks/src/gatling/scala/org/apache/polaris/benchmarks/actions/SetupActions.scala: ########## @@ -0,0 +1,132 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.polaris.benchmarks.actions + +import io.gatling.core.Predef._ +import io.gatling.core.structure.{ChainBuilder, ScenarioBuilder} +import org.apache.polaris.benchmarks.parameters.ConnectionParameters + +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} +import scala.concurrent.duration._ + +/** + * Actions for setting up authentication and managing shared state across benchmark simulations. + * This class encapsulates the common authentication patterns used across all simulations, + * eliminating code duplication and providing a consistent authentication mechanism. + * + * @param cp Connection parameters containing client credentials + * @param maxRetries Maximum number of retry attempts for authentication failures + * @param retryableHttpCodes HTTP status codes that should trigger a retry + */ +case class SetupActions( + cp: ConnectionParameters, + maxRetries: Int = 10, + retryableHttpCodes: Set[Int] = Set(500) +) { + + /** + * Shared access token reference that can be passed to all action classes. This token is + * automatically updated by the authentication scenarios. + */ + val accessToken: AtomicReference[String] = new AtomicReference() + + /** + * Internal flag to control the token refresh loop. Set to false to stop continuous token refresh + * scenarios. + */ + private val shouldRefreshToken: AtomicBoolean = new AtomicBoolean(true) + + /** + * Authentication actions instance that handles the actual OAuth token operations. This is private + * as all necessary functionality is exposed through SetupActions methods. + */ + private val authActions: AuthenticationActions = + AuthenticationActions(cp, accessToken, maxRetries, retryableHttpCodes) + + /** + * Continuously refreshes the OAuth token at a specified interval until stopped. This scenario + * runs in a loop controlled by the shouldRefreshToken flag. + * + * @param refreshInterval Duration between token refresh attempts (default: 1 minute) + * @return ScenarioBuilder that continuously refreshes the token + */ + def continuouslyRefreshOauthToken(refreshInterval: FiniteDuration = 1.minute): ScenarioBuilder = Review Comment: this parameter turned into a new `refresh-interval-seconds` config option under `auth` setting -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
