CORE-308 expand concurrent requests to sam (#4830)

DataBiosphere · Feb 27, 2025 · c1d60cd · c1d60cd
1 parent f0dd5f9
commit c1d60cd
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 9 deletions.
diff --git a/http/src/main/scala/org/broadinstitute/dsde/workbench/leonardo/config/Config.scala b/http/src/main/scala/org/broadinstitute/dsde/workbench/leonardo/config/Config.scala
@@ -385,7 +385,8 @@ object Config {
       Uri.unsafeFromString(config.as[String]("samServer")),
       config.getOrElse("petKeyCacheEnabled", true),
       config.getAs[FiniteDuration]("petKeyCacheExpiryTime").getOrElse(60 minutes),
-      config.getAs[Int]("petKeyCacheMaxSize").getOrElse(1000)
+      config.getAs[Int]("petKeyCacheMaxSize").getOrElse(1000),
+      config.getAs[Int]("maxConcurrentRequests").getOrElse(15)
     )
   }
 

diff --git a/http/src/main/scala/org/broadinstitute/dsde/workbench/leonardo/dao/HttpSamDAO.scala b/http/src/main/scala/org/broadinstitute/dsde/workbench/leonardo/dao/HttpSamDAO.scala
@@ -578,7 +578,8 @@ final case class ListResourceRolesItem[R](samResourceId: R, samRoles: Set[SamRol
 final case class HttpSamDaoConfig(samUri: Uri,
                                   petCacheEnabled: Boolean,
                                   petCacheExpiryTime: FiniteDuration,
-                                  petCacheMaxSize: Int
+                                  petCacheMaxSize: Int,
+                                  maxConcurrentRequests: Int
 )
 
 final case class UserEmailAndProject(userEmail: WorkbenchEmail, googleProject: GoogleProject)

diff --git a/.../main/scala/org/broadinstitute/dsde/workbench/leonardo/dao/sam/SamApiClientProvider.scala b/.../main/scala/org/broadinstitute/dsde/workbench/leonardo/dao/sam/SamApiClientProvider.scala
@@ -3,7 +3,7 @@ package org.broadinstitute.dsde.workbench.leonardo.dao.sam
 import cats.effect.Async
 import cats.mtl.Ask
 import cats.syntax.all._
-import okhttp3.Protocol
+import okhttp3.{Dispatcher, Protocol}
 import org.broadinstitute.dsde.workbench.client.sam.ApiClient
 import org.broadinstitute.dsde.workbench.client.sam.api.{AzureApi, GoogleApi, ResourcesApi, UsersApi}
 import org.broadinstitute.dsde.workbench.leonardo.AppContext
@@ -26,16 +26,24 @@ trait SamApiClientProvider[F[_]] {
   def azureApi(token: String)(implicit ev: Ask[F, AppContext]): F[AzureApi]
 }
 
-class HttpSamApiClientProvider[F[_]](samUrl: String)(implicit F: Async[F]) extends SamApiClientProvider[F] {
-  private val okHttpClient = new ApiClient().getHttpClient
+class HttpSamApiClientProvider[F[_]](samUrl: String, maxConcurrentRequests: Int)(implicit F: Async[F])
+    extends SamApiClientProvider[F] {
   private val timeout = 30 seconds
+  private val okHttpClient = {
+    val dispatcher = new Dispatcher()
+    dispatcher.setMaxRequests(maxConcurrentRequests)
+    dispatcher.setMaxRequestsPerHost(maxConcurrentRequests)
+    new ApiClient().getHttpClient.newBuilder
+      .readTimeout(timeout.toJava)
+      .protocols(Seq(Protocol.HTTP_1_1).asJava)
+      .dispatcher(dispatcher)
+      .build()
+  }
 
   private def getApiClient(token: String)(implicit ev: Ask[F, AppContext]): F[ApiClient] =
     for {
       ctx <- ev.ask
       okHttpClientBuilder = okHttpClient.newBuilder
-        .readTimeout(timeout.toJava)
-        .protocols(Seq(Protocol.HTTP_1_1).asJava)
       // TODO add otel interceptors
       //  See https://broadworkbench.atlassian.net/browse/IA-5052
       apiClient = new ApiClient(okHttpClientBuilder.build()).setBasePath(samUrl)

diff --git a/...n/scala/org/broadinstitute/dsde/workbench/leonardo/http/BaselineDependenciesBuilder.scala b/...n/scala/org/broadinstitute/dsde/workbench/leonardo/http/BaselineDependenciesBuilder.scala
@@ -120,7 +120,9 @@ class BaselineDependenciesBuilder {
                                                          applicationConfig
       )
 
-      samClientProvider = new HttpSamApiClientProvider(httpSamDaoConfig.samUri.renderString)
+      samClientProvider = new HttpSamApiClientProvider(httpSamDaoConfig.samUri.renderString,
+                                                       httpSamDaoConfig.maxConcurrentRequests
+      )
       samService = new SamServiceInterp(samClientProvider, cloudAuthTokenProvider)
 
       samDao <- buildHttpClient(sslContext, proxyResolver.resolveHttp4s, Some("leo_sam_client"), true).map(client =>

diff --git a/http/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/dao/HttpSamDAOSpec.scala b/http/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/dao/HttpSamDAOSpec.scala
@@ -32,7 +32,7 @@ import scala.util.control.NoStackTrace
 
 class HttpSamDAOSpec extends AnyFlatSpec with LeonardoTestSuite with BeforeAndAfterAll with MockitoSugar {
   val cloudAuthProvider = mock[CloudAuthTokenProvider[IO]]
-  val config = HttpSamDaoConfig(Uri.unsafeFromString("localhost"), false, 1 seconds, 10)
+  val config = HttpSamDaoConfig(Uri.unsafeFromString("localhost"), false, 1 seconds, 10, 15)
   implicit def unsafeLogger: Logger[IO] = Slf4jLogger.getLogger[IO]
   val underlyingPetTokenCache = Caffeine
     .newBuilder()