[DT-821]Possibly fix 401 and 500 error during dataset creation using retries

src/main/java/bio/terra/service/dataset/flight/create/CreateDatasetRegisterIngestServiceAccountStep.java

@@ -6,10 +6,14 @@
 import bio.terra.stairway.FlightMap;
 import bio.terra.stairway.Step;
 import bio.terra.stairway.StepResult;
+import bio.terra.stairway.StepStatus;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /** The step is only meant to be invoked for GCP backed datasets. */
 public class CreateDatasetRegisterIngestServiceAccountStep implements Step {
-
+  private static final Logger logger =
+      LoggerFactory.getLogger(CreateDatasetRegisterIngestServiceAccountStep.class);
   private final IamService iamService;
 
   public CreateDatasetRegisterIngestServiceAccountStep(IamService iamService) {
@@ -22,7 +26,17 @@ public StepResult doStep(FlightContext context) throws InterruptedException {
     String datasetServiceAccount =
         workingMap.get(DatasetWorkingMapKeys.SERVICE_ACCOUNT_EMAIL, String.class);
 
-    iamService.registerUser(datasetServiceAccount);
+    try {
+      iamService.registerUser(datasetServiceAccount);
+    } catch (Exception e) {
+      // catch transient errors from Sam, where the service account created in the previous step
+      // may not be ready to use yet
+      logger.warn(
+          String.format(
+              "Service account, %s, may not be ready to use yet. Retrying.", datasetServiceAccount),
+          e);
+      return new StepResult(StepStatus.STEP_RESULT_FAILURE_RETRY, e);
+    }
 
     return StepResult.getStepResultSuccess();
   }

src/main/java/bio/terra/service/dataset/flight/create/DatasetCreateFlight.java

@@ -102,7 +102,9 @@ public DatasetCreateFlight(FlightMap inputParameters, Object applicationContext)
       // Create the service account to use to ingest data and register it in Terra
       if (datasetRequest.isDedicatedIngestServiceAccount()) {
         addStep(new CreateDatasetCreateIngestServiceAccountStep(resourceService, datasetRequest));
-        addStep(new CreateDatasetRegisterIngestServiceAccountStep(iamService));
+        addStep(
+            new CreateDatasetRegisterIngestServiceAccountStep(iamService),
+            getDefaultExponentialBackoffRetryRule());
       }
     }
 

src/test/java/bio/terra/service/dataset/flight/create/CreateDatasetRegisterIngestServiceAccountStepTest.java

@@ -0,0 +1,60 @@
+package bio.terra.service.dataset.flight.create;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import bio.terra.common.category.Unit;
+import bio.terra.service.auth.iam.IamService;
+import bio.terra.service.auth.iam.exception.IamInternalServerErrorException;
+import bio.terra.service.auth.iam.exception.IamUnauthorizedException;
+import bio.terra.service.dataset.flight.DatasetWorkingMapKeys;
+import bio.terra.stairway.FlightContext;
+import bio.terra.stairway.FlightMap;
+import bio.terra.stairway.StepStatus;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+@Tag(Unit.TAG)
+@ExtendWith(MockitoExtension.class)
+class CreateDatasetRegisterIngestServiceAccountStepTest {
+  @Mock private IamService iamService;
+  @Mock private FlightContext flightContext;
+  private CreateDatasetRegisterIngestServiceAccountStep step;
+
+  @BeforeEach
+  void setup() {
+    step = new CreateDatasetRegisterIngestServiceAccountStep(iamService);
+    FlightMap workingMap = new FlightMap();
+    workingMap.put(DatasetWorkingMapKeys.SERVICE_ACCOUNT_EMAIL, "email");
+    when(flightContext.getWorkingMap()).thenReturn(workingMap);
+  }
+
+  @Test
+  void doStep() throws InterruptedException {
+    assertThat(step.doStep(flightContext).getStepStatus(), equalTo(StepStatus.STEP_RESULT_SUCCESS));
+    verify(iamService).registerUser("email");
+  }
+
+  @Test
+  void doStep_Retry401() throws InterruptedException {
+    doThrow(new IamUnauthorizedException("Unauthorized")).when(iamService).registerUser("email");
+    assertThat(
+        step.doStep(flightContext).getStepStatus(), equalTo(StepStatus.STEP_RESULT_FAILURE_RETRY));
+  }
+
+  @Test
+  void doStep_500() throws InterruptedException {
+    doThrow(new IamInternalServerErrorException("Internal Server Error"))
+        .when(iamService)
+        .registerUser("email");
+    assertThat(
+        step.doStep(flightContext).getStepStatus(), equalTo(StepStatus.STEP_RESULT_FAILURE_RETRY));
+  }
+}