Exponentially retry 409s from GoogleIamDAO policy changes, to support concurrent cluster creations in the same project (#737)

Author: rtitle

src/main/scala/org/broadinstitute/dsde/workbench/leonardo/monitor/ClusterMonitorActor.scala

@@ -3,10 +3,11 @@ package org.broadinstitute.dsde.workbench.leonardo.monitor
 import java.time.Instant
 
 import akka.actor.Status.Failure
-import akka.actor.{Actor, Props}
+import akka.actor.{Actor, ActorSystem, Props}
 import akka.pattern.pipe
 import cats.data.OptionT
 import cats.implicits._
+import com.google.api.client.googleapis.json.GoogleJsonResponseException
 import com.typesafe.scalalogging.LazyLogging
 import io.grpc.Status.Code
 import org.broadinstitute.dsde.workbench.google.{GoogleIamDAO, GoogleStorageDAO}
@@ -19,7 +20,7 @@ import org.broadinstitute.dsde.workbench.leonardo.model.google.ClusterStatus._
 import org.broadinstitute.dsde.workbench.leonardo.model.google.{ClusterStatus, IP, _}
 import org.broadinstitute.dsde.workbench.leonardo.monitor.ClusterMonitorActor._
 import org.broadinstitute.dsde.workbench.leonardo.monitor.ClusterMonitorSupervisor.{ClusterDeleted, ClusterSupervisorMessage, RemoveFromList}
-import org.broadinstitute.dsde.workbench.util.addJitter
+import org.broadinstitute.dsde.workbench.util.{Retry, addJitter}
 import slick.dbio.DBIOAction
 
 import scala.collection.immutable.Set
@@ -63,9 +64,12 @@ class ClusterMonitorActor(val cluster: Cluster,
                           val googleStorageDAO: GoogleStorageDAO,
                           val dbRef: DbReference,
                           val authProvider: LeoAuthProvider,
-                          val jupyterProxyDAO: JupyterDAO) extends Actor with LazyLogging {
+                          val jupyterProxyDAO: JupyterDAO) extends Actor with LazyLogging with Retry {
   import context._
 
+  // the Retry trait needs a reference to the ActorSystem
+  override val system = context.system
+
   override def preStart(): Unit = {
     super.preStart()
     scheduleInitialMonitorPass
@@ -342,6 +346,13 @@ class ClusterMonitorActor(val cluster: Cluster,
     transformed.value
   }
 
+  private def whenGoogle409(throwable: Throwable): Boolean = {
+    throwable match {
+      case t: GoogleJsonResponseException => t.getStatusCode == 409
+      case _ => false
+    }
+  }
+
   private def removeIamRolesForUser: Future[Unit] = {
     // Remove the Dataproc Worker IAM role for the cluster service account
     cluster.serviceAccountInfo.clusterServiceAccount match {
@@ -356,7 +367,11 @@ class ClusterMonitorActor(val cluster: Cluster,
           if (count > 0) {
             Future.successful(())
           } else {
-            googleIamDAO.removeIamRolesForUser(cluster.googleProject, serviceAccountEmail, Set("roles/dataproc.worker"))
+            // Retry 409s with exponential backoff. This can happen if concurrent policy updates are made in the same project.
+            // Google recommends a retry in this case.
+            retryExponentially(whenGoogle409, s"IAM policy change failed for Google project '${cluster.googleProject}'") { () =>
+              googleIamDAO.removeIamRolesForUser(cluster.googleProject, serviceAccountEmail, Set("roles/dataproc.worker"))
+            }
           }
         }
     }

src/main/scala/org/broadinstitute/dsde/workbench/leonardo/service/LeonardoService.scala

@@ -1,6 +1,6 @@
 package org.broadinstitute.dsde.workbench.leonardo.service
 
-import java.io.File
+import java.io.{File, IOException}
 import java.time.Instant
 
 import akka.actor.ActorSystem
@@ -686,15 +686,32 @@ class LeonardoService(protected val dataprocConfig: DataprocConfig,
     tea.value.void
   }
 
+  private def whenGoogle409(throwable: Throwable): Boolean = {
+    throwable match {
+      case t: GoogleJsonResponseException => t.getStatusCode == 409
+      case _ => false
+    }
+  }
+
   private[service] def addDataprocWorkerRoleToServiceAccount(googleProject: GoogleProject, serviceAccountOpt: Option[WorkbenchEmail]): Future[Unit] = {
     serviceAccountOpt.map { serviceAccountEmail =>
-      googleIamDAO.addIamRolesForUser(googleProject, serviceAccountEmail, Set("roles/dataproc.worker"))
+      // Retry 409s with exponential backoff. This can happen if concurrent policy updates are made in the same project.
+      // Google recommends a retry in this case.
+      val iamFuture: Future[Unit] = retryExponentially(whenGoogle409, s"IAM policy change failed for Google project '$googleProject'") { () =>
+        googleIamDAO.addIamRolesForUser(googleProject, serviceAccountEmail, Set("roles/dataproc.worker"))
+      }
+      iamFuture
     } getOrElse Future.successful(())
   }
 
   private[service] def removeDataprocWorkerRoleFromServiceAccount(googleProject: GoogleProject, serviceAccountOpt: Option[WorkbenchEmail]): Future[Unit] = {
     serviceAccountOpt.map { serviceAccountEmail =>
-      googleIamDAO.removeIamRolesForUser(googleProject, serviceAccountEmail, Set("roles/dataproc.worker"))
+      // Retry 409s with exponential backoff. This can happen if concurrent policy updates are made in the same project.
+      // Google recommends a retry in this case.
+      val iamFuture: Future[Unit] = retryExponentially(whenGoogle409, s"IAM policy change failed for Google project '$googleProject'") { () =>
+        googleIamDAO.removeIamRolesForUser(googleProject, serviceAccountEmail, Set("roles/dataproc.worker"))
+      }
+      iamFuture
     } getOrElse Future.successful(())
   }
 

src/test/scala/org/broadinstitute/dsde/workbench/leonardo/service/LeonardoServiceSpec.scala

@@ -9,6 +9,7 @@ import akka.http.scaladsl.model.headers.OAuth2BearerToken
 import akka.testkit.TestKit
 import com.google.api.client.googleapis.testing.json.GoogleJsonResponseExceptionFactoryTesting
 import com.google.api.client.testing.json.MockJsonFactory
+import com.typesafe.scalalogging.LazyLogging
 import org.broadinstitute.dsde.workbench.google.GoogleStorageDAO
 import org.broadinstitute.dsde.workbench.google.mock.{MockGoogleDataprocDAO, MockGoogleIamDAO, MockGoogleStorageDAO}
 import org.broadinstitute.dsde.workbench.leonardo.CommonTestData
@@ -25,19 +26,21 @@ import org.broadinstitute.dsde.workbench.leonardo.monitor.NoopActor
 import org.broadinstitute.dsde.workbench.leonardo.util.BucketHelper
 import org.broadinstitute.dsde.workbench.model._
 import org.broadinstitute.dsde.workbench.model.google._
+import org.broadinstitute.dsde.workbench.util.Retry
 import org.mockito.ArgumentMatchers.{any, eq => mockitoEq}
 import org.mockito.Mockito.{never, verify, _}
 import org.scalatest._
 import org.scalatest.concurrent.Eventually.eventually
 import org.scalatest.concurrent.ScalaFutures
+import org.scalatest.time.{Minutes, Span}
 import spray.json._
 
 import scala.concurrent.duration._
 import scala.concurrent.{Await, ExecutionContext, Future}
 
 class LeonardoServiceSpec extends TestKit(ActorSystem("leonardotest")) with FlatSpecLike with Matchers
   with BeforeAndAfter with BeforeAndAfterAll with TestComponent with ScalaFutures
-  with OptionValues with CommonTestData with LeoComponent {
+  with OptionValues with CommonTestData with LeoComponent with Retry with LazyLogging {
 
   private var gdDAO: MockGoogleDataprocDAO = _
   private var computeDAO: MockGoogleComputeDAO = _
@@ -975,8 +978,9 @@ class LeonardoServiceSpec extends TestKit(ActorSystem("leonardotest")) with Flat
       new MockGoogleStorageDAO
     }
 
-    //we meed to use a special version of the MockGoogleDataprocDAO to simulate an error during the call to resizeCluster
-    leo = new LeonardoService(dataprocConfig, clusterFilesConfig, clusterResourcesConfig, clusterDefaultsConfig, proxyConfig, swaggerConfig, autoFreezeConfig, mockGoogleDataprocDAO, computeDAO, new ErroredMockGoogleIamDAO, storageDAO, mockPetGoogleDAO, DbSingleton.ref, authProvider, serviceAccountProvider, whitelist, bucketHelper, contentSecurityPolicy)
+    //we meed to use a special version of the MockGoogleIamDAO to simulate an error when adding IAM roles
+    val iamDAO = new ErroredMockGoogleIamDAO
+    leo = new LeonardoService(dataprocConfig, clusterFilesConfig, clusterResourcesConfig, clusterDefaultsConfig, proxyConfig, swaggerConfig, autoFreezeConfig, mockGoogleDataprocDAO, computeDAO, iamDAO, storageDAO, mockPetGoogleDAO, DbSingleton.ref, authProvider, serviceAccountProvider, whitelist, bucketHelper, contentSecurityPolicy)
 
     // create the cluster
     val clusterCreateResponse =
@@ -985,6 +989,30 @@ class LeonardoServiceSpec extends TestKit(ActorSystem("leonardotest")) with Flat
     eventually {
       dbFutureValue { _.clusterQuery.getClusterStatus(clusterCreateResponse.id) } shouldBe Some(ClusterStatus.Error)
     }
+
+    // IAM call should not have been retried
+    iamDAO.invocationCount shouldBe 1
+  }
+
+  it should "retry 409 errors when adding dataproc worker role" in isolatedDbTest {
+    val mockPetGoogleDAO: String => GoogleStorageDAO = _ => {
+      new MockGoogleStorageDAO
+    }
+
+    //we meed to use a special version of the MockGoogleIamDAO to simulate a conflict when adding IAM roles
+    val iamDAO = new ErroredMockGoogleIamDAO(409)
+    leo = new LeonardoService(dataprocConfig, clusterFilesConfig, clusterResourcesConfig, clusterDefaultsConfig, proxyConfig, swaggerConfig, autoFreezeConfig, mockGoogleDataprocDAO, computeDAO, iamDAO, storageDAO, mockPetGoogleDAO, DbSingleton.ref, authProvider, serviceAccountProvider, whitelist, bucketHelper, contentSecurityPolicy)
+
+    // create the cluster
+    val clusterCreateResponse =
+      leo.processClusterCreationRequest(userInfo, project, name1, testClusterRequest).futureValue
+
+    eventually(timeout(Span(5, Minutes))) {
+      dbFutureValue { _.clusterQuery.getClusterStatus(clusterCreateResponse.id) } shouldBe Some(ClusterStatus.Error)
+    }
+
+    // IAM call should have been retried exponentially
+    iamDAO.invocationCount shouldBe exponentialBackOffIntervals.size + 1
   }
 
   it should "update the autopause threshold for a cluster" in isolatedDbTest {
@@ -1158,10 +1186,12 @@ class LeonardoServiceSpec extends TestKit(ActorSystem("leonardotest")) with Flat
     }
   }
 
-  private class ErroredMockGoogleIamDAO extends MockGoogleIamDAO {
+  private class ErroredMockGoogleIamDAO(statusCode: Int = 400) extends MockGoogleIamDAO {
+    var invocationCount = 0
     override def addIamRolesForUser(iamProject: GoogleProject, email: WorkbenchEmail, rolesToAdd: Set[String]): Future[Unit] = {
+      invocationCount += 1
       val jsonFactory = new MockJsonFactory
-      val testException = GoogleJsonResponseExceptionFactoryTesting.newMock(jsonFactory, 400, "oh no i have failed")
+      val testException = GoogleJsonResponseExceptionFactoryTesting.newMock(jsonFactory, statusCode, "oh no i have failed")
 
       Future.failed(testException)
     }