diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datastream/DatastreamResourceManager.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datastream/DatastreamResourceManager.java index 8fc52ac4bf..ade6779a8b 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datastream/DatastreamResourceManager.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datastream/DatastreamResourceManager.java @@ -69,6 +69,7 @@ import java.util.Set; import java.util.concurrent.CancellationException; import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicBoolean; import org.apache.arrow.util.VisibleForTesting; import org.apache.beam.it.common.ResourceManager; import org.apache.beam.it.common.utils.ExceptionUtils; @@ -376,7 +377,9 @@ public synchronized DestinationConfig buildGCSDestinationConfig( DestinationConfig.newBuilder().setDestinationConnectionProfile(connectionProfile.getName()); GcsDestinationConfig.Builder gcsDestinationConfigBuilder = - GcsDestinationConfig.newBuilder().setPath(path); + GcsDestinationConfig.newBuilder() + .setPath(path) + .setFileRotationInterval(Duration.newBuilder().setSeconds(60).build()); if (destinationOutputFormat == DestinationOutputFormat.AVRO_FILE_FORMAT) { gcsDestinationConfigBuilder.setAvroFileFormat(AvroFileFormat.getDefaultInstance()); @@ -637,11 +640,12 @@ public synchronized Stream pauseStream(Stream stream) { } public synchronized void cleanupAll() { - LOG.info("Cleaning up Datastream resource manager."); - boolean producedError = false; + LOG.info("Cleaning up Datastream resource manager at {}.", java.time.Instant.now()); + AtomicBoolean producedError = new AtomicBoolean(false); for (String stream : createdStreamIds) { try { + LOG.info("Starting deleteStreamAsync get for {} at {}", stream, java.time.Instant.now()); Failsafe.with(retryOnException()) .get( () -> @@ -651,6 +655,7 @@ public synchronized void cleanupAll() { .setName(StreamName.format(projectId, location, stream)) .build()) .get()); + LOG.info("Finished deleteStreamAsync get for {} at {}", stream, java.time.Instant.now()); } catch (Exception e) { if (ExceptionUtils.containsType(e, NotFoundException.class) || ExceptionUtils.containsMessage(e, "NOT_FOUND")) { @@ -658,48 +663,62 @@ public synchronized void cleanupAll() { "Stream {} not found in project {}. Assuming already deleted.", stream, projectId); } else { LOG.error("Failed to delete stream {}.", stream, e); - producedError = true; + producedError.set(true); } } } - LOG.info("Successfully deleted stream(s). "); - - for (String connectionProfile : createdConnectionProfileIds) { - try { - Failsafe.with(retryOnException()) - .get( - () -> - datastreamClient - .deleteConnectionProfileAsync( - DeleteConnectionProfileRequest.newBuilder() - .setName( - ConnectionProfileName.format( - projectId, location, connectionProfile)) - .build()) - .get()); - } catch (Exception e) { - if (ExceptionUtils.containsType(e, NotFoundException.class) - || ExceptionUtils.containsMessage(e, "NOT_FOUND")) { - LOG.warn( - "Connection Profile {} not found in project {}. Assuming already deleted.", - connectionProfile, - projectId); - } else { - LOG.error("Failed to delete connection profile {}.", connectionProfile, e); - producedError = true; - } - } - } - LOG.info("Successfully deleted connection profile(s). "); + LOG.info("Successfully deleted stream(s) at {}. ", java.time.Instant.now()); + + createdConnectionProfileIds.parallelStream() + .forEach( + connectionProfile -> { + try { + LOG.info( + "Starting deleteConnectionProfileAsync get for {} at {}", + connectionProfile, + java.time.Instant.now()); + Failsafe.with(retryOnException()) + .get( + () -> + datastreamClient + .deleteConnectionProfileAsync( + DeleteConnectionProfileRequest.newBuilder() + .setName( + ConnectionProfileName.format( + projectId, location, connectionProfile)) + .build()) + .get()); + LOG.info( + "Finished deleteConnectionProfileAsync get for {} at {}", + connectionProfile, + java.time.Instant.now()); + } catch (Exception e) { + if (ExceptionUtils.containsType(e, NotFoundException.class) + || ExceptionUtils.containsMessage(e, "NOT_FOUND")) { + LOG.warn( + "Connection Profile {} not found in project {}. Assuming already deleted.", + connectionProfile, + projectId); + } else { + LOG.error( + "Failed to delete connection profile {} at {}.", + connectionProfile, + java.time.Instant.now(), + e); + producedError.set(true); + } + } + }); + LOG.info("Successfully deleted connection profile(s) at {}. ", java.time.Instant.now()); try { datastreamClient.close(); } catch (Exception e) { LOG.error("Failed to close datastream client. "); - producedError = true; + producedError.set(true); } - if (producedError) { + if (producedError.get()) { throw new DatastreamResourceManagerException( "Failed to delete resources. Check above for errors."); } diff --git a/v2/datastream-to-bigquery/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamToBigQueryIT.java b/v2/datastream-to-bigquery/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamToBigQueryIT.java index 5bc48792d4..0874fd0411 100644 --- a/v2/datastream-to-bigquery/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamToBigQueryIT.java +++ b/v2/datastream-to-bigquery/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamToBigQueryIT.java @@ -67,6 +67,7 @@ import org.apache.commons.lang3.RandomStringUtils; import org.checkerframework.checker.nullness.qual.NonNull; import org.junit.After; +import org.junit.AfterClass; import org.junit.Before; import org.junit.Ignore; import org.junit.Rule; @@ -103,16 +104,19 @@ enum JDBCType { private CloudSqlResourceManager cloudSqlResourceManager; private PubsubResourceManager pubsubResourceManager; - private DatastreamResourceManager datastreamResourceManager; + private static DatastreamResourceManager datastreamResourceManager; private BigQueryResourceManager bigQueryResourceManager; @Before public void setUp() throws IOException { - datastreamResourceManager = - DatastreamResourceManager.builder(testName, PROJECT, REGION) - .setCredentialsProvider(credentialsProvider) - .setPrivateConnectivity("datastream-connect-2") - .build(); + if (datastreamResourceManager == null) { + datastreamResourceManager = + DatastreamResourceManager.builder( + DataStreamToBigQueryIT.class.getSimpleName(), PROJECT, REGION) + .setCredentialsProvider(credentialsProvider) + .setPrivateConnectivity("datastream-connect-2") + .build(); + } bigQueryResourceManager = BigQueryResourceManager.builder(testName, PROJECT, credentials).build(); @@ -143,10 +147,12 @@ public void setUp() throws IOException { @After public void cleanUp() { ResourceManagerUtils.cleanResources( - cloudSqlResourceManager, - pubsubResourceManager, - datastreamResourceManager, - bigQueryResourceManager); + cloudSqlResourceManager, pubsubResourceManager, bigQueryResourceManager); + } + + @AfterClass + public static void tearDownClass() { + ResourceManagerUtils.cleanResources(datastreamResourceManager); } @Test