Dataflow runner does not respect ParDo's lifecycle on case of exceptions

[kennknowles]

The lifecycle of the DoFn is not respected in case of exception in any of the lifecycle methods after setup.

Imported from Jira BEAM-3245. Original Jira may contain additional context.
Reported by: iemejia.

[Abacn]

Found that tearDown is never called in Dataflow runner v1 or v2, even pipeline finished normally

Just use a simple DoFn:

    pipeline
        .apply(GenerateSequence.from(0).to(1000))
        .apply(ParDo.of(new SomeDoFn()));

where SomeDoFn is

static class SomeDoFn extends DoFn<Long, Long> {

    protected transient Integer id;

    @Setup
    public void setup() {
      Integer idOld = id;
      id = ThreadLocalRandom.current().nextInt(0, 1000);
      LOG.info("@Setup {}, {} at {}", idOld, id, Integer.toHexString(System.identityHashCode(this)));
    }

    @StartBundle
    public void startBundle() {
      LOG.info("@StartBundle {} at {}", id, Integer.toHexString(System.identityHashCode(this)));
    }

    @ProcessElement
    public void process(@Element Long input, OutputReceiver<Long> receiver) {
      // LOG.info("@ProcessElement for {} at bundle {}", input, id);
      receiver.output(input);
    }

    @FinishBundle
    public void finishBundle() throws Exception {
      LOG.info("@FinishBundle {} at {}", id, Integer.toHexString(System.identityHashCode(this)));
    }

    @Teardown
    public void teardown() {
      LOG.info("@Teardown {} at {}", id, Integer.toHexString(System.identityHashCode(this)));
    }

    @Override
    protected void finalize() throws Throwable {
      LOG.info("finalize");
    }
  }

seen log in setup, start/finishBundle, but not teardown.

[Abacn]

Checked that for Dataflow legacy runner, teardown is called after a DoFn throw error in processElement or finishBundle, but not Dataflow runner v2.

update: runner v2 case partly fixed by #32522

[scwhittle]

One previously identified issue was that teardown may not be invoked on a DoFn if there is an exception calling setup on a different fused DoFn.

However I believe that if the dataflow runner abandons a MapTaskExecutor for some reason not related to an exception in user processing (cache timeout or exception during backend commit) that we may not be invoking the teardown method. This can lead to leaking resources if they are created in setup and require teardown to close them. An example of this is the KafkaProducer in KafkaWriter here.

[Abacn]

One previously identified issue was that teardown may not be invoked on a DoFn if there is an exception calling setup on a different fused DoFn.

Yes, link here: https://github.com/apache/beam/pull/32522/files#r1777301355

However I believe that if the dataflow runner abandons a MapTaskExecutor for some reason not related to an exception in user processing (cache timeout or exception during backend commit) that we may not be invoking the teardown method. This can lead to leaking resources if they are created in setup and require teardown to close them. An example of this is the KafkaProducer in KafkaWriter here.

The link not working, pasted a perma-link here for reference:

beam/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriter.java

Line 57 in c22665c

producer = new KafkaProducer<>(producerConfig);

[scwhittle]

I put together a PR to fix and will be working next on a test that verifies this is the case in Dataflow Streaming runner but it certainly appears to be so from looking at the code.

[scwhittle]

@Abacn Are there remaining issues here that are known or can I close this?