diff --git a/automl/snippets/src/test/java/com/example/automl/ListOperationStatusTest.java b/automl/snippets/src/test/java/com/example/automl/ListOperationStatusTest.java index 2b60834aa61..432365e8949 100644 --- a/automl/snippets/src/test/java/com/example/automl/ListOperationStatusTest.java +++ b/automl/snippets/src/test/java/com/example/automl/ListOperationStatusTest.java @@ -19,9 +19,19 @@ import static com.google.common.truth.Truth.assertThat; import static junit.framework.TestCase.assertNotNull; +import com.google.api.client.util.ExponentialBackOff; +import com.google.api.gax.rpc.ResourceExhaustedException; +import com.google.cloud.automl.v1.AutoMlClient; +import com.google.cloud.automl.v1.LocationName; +import com.google.longrunning.ListOperationsRequest; +import com.google.longrunning.Operation; +import com.google.longrunning.OperationsClient; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.PrintStream; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; import org.junit.After; import org.junit.Before; import org.junit.BeforeClass; @@ -31,7 +41,7 @@ @RunWith(JUnit4.class) public class ListOperationStatusTest { - private static final String PROJECT_ID = System.getenv("AUTOML_PROJECT_ID"); + private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT"); private ByteArrayOutputStream bout; private PrintStream out; private PrintStream originalPrintStream; @@ -49,11 +59,64 @@ public static void checkRequirements() { } @Before - public void setUp() { + public void setUp() throws IOException, InterruptedException { bout = new ByteArrayOutputStream(); out = new PrintStream(bout); originalPrintStream = System.out; System.setOut(out); + + // if the LRO status count more than 300, delete half of operations. + try (AutoMlClient client = AutoMlClient.create()) { + OperationsClient operationsClient = client.getOperationsClient(); + LocationName projectLocation = LocationName.of(PROJECT_ID, "us-central1"); + ListOperationsRequest listRequest = + ListOperationsRequest.newBuilder().setName(projectLocation.toString()).build(); + List operationFullPathsToBeDeleted = new ArrayList<>(); + for (Operation operation : operationsClient.listOperations(listRequest).iterateAll()) { + // collect unused operation into the list. + // Filter: deleting already done operations. + if (operation.getDone() && !operation.hasError()) { + operationFullPathsToBeDeleted.add(operation.getName()); + } + } + + if (operationFullPathsToBeDeleted.size() > 300) { + System.out.println("Cleaning up..."); + + + for (String operationFullPath : + operationFullPathsToBeDeleted.subList(0, operationFullPathsToBeDeleted.size() / 2)) { + // retry_interval * (random value in range [1 - rand_factor, 1 + rand_factor]) + ExponentialBackOff exponentialBackOff = new ExponentialBackOff.Builder() + .setInitialIntervalMillis(60000) + .setMaxElapsedTimeMillis(300000) + .setRandomizationFactor(0.5) + .setMultiplier(1.1) + .setMaxIntervalMillis(80000) + .build(); + + // delete unused operations. + try { + operationsClient.deleteOperation(operationFullPath); + } catch (ResourceExhaustedException ex) { + // exponential back off and retry. + long backOffInMillis = exponentialBackOff.nextBackOffMillis(); + System.out.printf("Backing off for %d milliseconds " + + "due to Resource exhaustion.\n", backOffInMillis); + if (backOffInMillis < 0) { + break; + } + System.out.println("Backing off" + backOffInMillis); + TimeUnit.MILLISECONDS.sleep(backOffInMillis); + } catch (Exception ex) { + throw ex; + } + } + } else { + // Clear the list since we wont anything with the list. + operationFullPathsToBeDeleted.clear(); + } + } } @After