Correcting Dataflow jobs resource leaks (GoogleCloudPlatform#3381)

Author: darshan-sj

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/TemplateTestBase.java

@@ -543,26 +543,10 @@ protected PipelineLauncher buildLauncher(
    * jobs getting leaked.
    */
   protected LaunchInfo launchTemplate(LaunchConfig.Builder options) throws IOException {
-    return this.launchTemplate(options, true, this.template);
+    return this.launchTemplate(options, this.template);
   }
 
-  /**
-   * Launch the template with the given options and configuration for hook.
-   *
-   * @param options Options to use for launch.
-   * @param setupShutdownHook Whether should setup a hook to cancel the job upon VM termination.
-   *     This is useful to teardown resources if the VM/test terminates unexpectedly.
-   * @return Job details.
-   * @throws IOException Thrown when {@link PipelineLauncher#launch(String, String, LaunchConfig)}
-   *     fails.
-   */
-  protected LaunchInfo launchTemplate(LaunchConfig.Builder options, boolean setupShutdownHook)
-      throws IOException {
-    return this.launchTemplate(options, setupShutdownHook, this.template);
-  }
-
-  protected LaunchInfo launchTemplate(
-      LaunchConfig.Builder options, boolean setupShutdownHook, Template templateMetadata)
+  protected LaunchInfo launchTemplate(LaunchConfig.Builder options, Template templateMetadata)
       throws IOException {
 
     boolean flex =
@@ -638,8 +622,8 @@ protected LaunchInfo launchTemplate(
 
     LaunchInfo launchInfo = pipelineLauncher.launch(PROJECT, REGION, options.build());
 
-    // if the launch succeeded and setupShutdownHook is enabled, setup a thread to cancel job
-    if (setupShutdownHook && launchInfo.jobId() != null && !launchInfo.jobId().isEmpty()) {
+    // if the launch succeeded, setup a thread to cancel job
+    if (launchInfo.jobId() != null && !launchInfo.jobId().isEmpty()) {
       Runtime.getRuntime()
           .addShutdownHook(new Thread(new CancelJobShutdownHook(pipelineLauncher, launchInfo)));
     }

v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamToSpannerITBase.java

@@ -417,7 +417,7 @@ protected LaunchInfo launchDataflowJob(
 
     // Run
     LOG.info("Launching Dataflow job with parameters: {}", params);
-    LaunchInfo jobInfo = launchTemplate(options, false);
+    LaunchInfo jobInfo = launchTemplate(options);
     assertThatPipeline(jobInfo).isRunning();
 
     return jobInfo;

v2/gcs-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/GCSToSourceDbInterleaveIT.java

@@ -382,7 +382,7 @@ private void launchWriterDataflowJob() throws IOException {
     LaunchConfig.Builder options = LaunchConfig.builder(jobName, specPath);
     options.setParameters(params);
     // Run
-    writerJobInfo = launchTemplate(options, false);
+    writerJobInfo = launchTemplate(options);
   }
 
   private void launchReaderDataflowJob() throws IOException {

v2/gcs-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/GCSToSourceDbInterleaveMultiShardIT.java

@@ -419,7 +419,7 @@ private void launchWriterDataflowJob() throws IOException {
     LaunchConfig.Builder options = LaunchConfig.builder(jobName, specPath);
     options.setParameters(params);
     // Run
-    writerJobInfo = launchTemplate(options, false);
+    writerJobInfo = launchTemplate(options);
   }
 
   private void launchReaderDataflowJob() throws IOException {

v2/gcs-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/GCSToSourceDbWithReaderIT.java

@@ -212,7 +212,7 @@ private void launchWriterDataflowJob() throws IOException {
     LaunchConfig.Builder options = LaunchConfig.builder(jobName, specPath);
     options.setParameters(params);
     // Run
-    writerJobInfo = launchTemplate(options, false);
+    writerJobInfo = launchTemplate(options);
   }
 
   private void launchReaderDataflowJob() throws IOException {

v2/gcs-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/GCSToSourceDbWithoutReaderIT.java

@@ -279,7 +279,7 @@ private void launchWriterDataflowJob(CustomTransformation customTransformation)
     LaunchConfig.Builder options = LaunchConfig.builder(jobName, specPath);
     options.setParameters(params);
     // Run
-    jobInfo = launchTemplate(options, false);
+    jobInfo = launchTemplate(options);
   }
 
   private void createAndUploadShardConfigToGcs(

v2/gcs-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/TimezoneIT.java

@@ -244,7 +244,7 @@ private void launchWriterDataflowJob() throws IOException {
     LaunchConfig.Builder options = LaunchConfig.builder(jobName, specPath);
     options.setParameters(params);
     // Run
-    writerJobInfo = launchTemplate(options, false);
+    writerJobInfo = launchTemplate(options);
   }
 
   private void launchReaderDataflowJob() throws IOException {

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java

@@ -274,7 +274,7 @@ protected PipelineLauncher.LaunchInfo launchDataflowJob(
     options.addEnvironment("numWorkers", 2);
     options.addEnvironment("ipConfiguration", "WORKER_IP_PRIVATE");
     // Run
-    PipelineLauncher.LaunchInfo jobInfo = launchTemplate(options, false);
+    PipelineLauncher.LaunchInfo jobInfo = launchTemplate(options);
     assertThatPipeline(jobInfo).isRunning();
 
     return jobInfo;

v2/spanner-change-streams-to-sharded-file-sink/src/test/java/com/google/cloud/teleport/v2/templates/SpannerChangeStreamToGcsITBase.java

@@ -176,7 +176,7 @@ public PipelineLauncher.LaunchInfo launchReaderDataflowJob(
     options.setParameters(params);
     options.addEnvironment("additionalExperiments", Collections.singletonList("use_runner_v2"));
     // Run
-    PipelineLauncher.LaunchInfo jobInfo = launchTemplate(options, false);
+    PipelineLauncher.LaunchInfo jobInfo = launchTemplate(options);
     assertThatPipeline(jobInfo).isRunning();
     return jobInfo;
   }

v2/spanner-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/SpannerToSourceDbITBase.java

@@ -275,7 +275,7 @@ public PipelineLauncher.LaunchInfo launchDataflowJob(
     options.addEnvironment("additionalExperiments", Collections.singletonList("use_runner_v2"));
     options.addEnvironment("ipConfiguration", "WORKER_IP_PRIVATE");
     // Run
-    PipelineLauncher.LaunchInfo jobInfo = launchTemplate(options, false);
+    PipelineLauncher.LaunchInfo jobInfo = launchTemplate(options);
     assertThatPipeline(jobInfo).isRunning();
     return jobInfo;
   }