expose more context in DoFns, fix spotify#3579 (spotify#3583)

- `BundleFinalizer` is not supported by `SimpleDoFnRunner` in `DirectRunner`

Author: nevillelyh

scio-core/src/main/java/com/spotify/scio/transforms/BaseAsyncDoFn.java

@@ -43,14 +43,14 @@ public abstract class BaseAsyncDoFn<InputT, OutputT, ResourceT, FutureT>
   private final ConcurrentLinkedQueue<Throwable> errors = new ConcurrentLinkedQueue<>();
 
   @StartBundle
-  public void startBundle() {
+  public void startBundle(StartBundleContext context) {
     futures.clear();
     results.clear();
     errors.clear();
   }
 
   @FinishBundle
-  public void finishBundle(FinishBundleContext c) {
+  public void finishBundle(FinishBundleContext context) {
     if (!futures.isEmpty()) {
       try {
         waitForFutures(futures.values());
@@ -63,7 +63,7 @@ public void finishBundle(FinishBundleContext c) {
         throw new RuntimeException("Failed to process futures", e);
       }
     }
-    flush(c);
+    flush(context);
   }
 
   @ProcessElement

scio-core/src/main/java/com/spotify/scio/transforms/BaseAsyncLookupDoFn.java

@@ -111,7 +111,7 @@ public void setup() {
   }
 
   @StartBundle
-  public void startBundle() {
+  public void startBundle(StartBundleContext context) {
     futures.clear();
     results.clear();
     requestCount = 0;
@@ -178,7 +178,7 @@ public void processElement(ProcessContext c, BoundedWindow window) {
   }
 
   @FinishBundle
-  public void finishBundle(FinishBundleContext c) {
+  public void finishBundle(FinishBundleContext context) {
     if (!futures.isEmpty()) {
       try {
         // Block until all pending futures are complete
@@ -191,7 +191,7 @@ public void finishBundle(FinishBundleContext c) {
         LOG.error("Failed to process futures", e);
       }
     }
-    flush(r -> c.output(KV.of(r.input, r.output), r.timestamp, r.window));
+    flush(r -> context.output(KV.of(r.input, r.output), r.timestamp, r.window));
 
     // Make sure all requests are processed
     Preconditions.checkState(

scio-core/src/main/java/com/spotify/scio/transforms/FileDownloadDoFn.java

@@ -75,7 +75,7 @@ public FileDownloadDoFn(
   }
 
   @StartBundle
-  public void startBundle() {
+  public void startBundle(StartBundleContext context) {
     this.batch.clear();
   }
 
@@ -88,8 +88,8 @@ public void processElement(ProcessContext c, BoundedWindow window) {
   }
 
   @FinishBundle
-  public void finishBundle(FinishBundleContext c) {
-    processBatch(c);
+  public void finishBundle(FinishBundleContext context) {
+    processBatch(context);
   }
 
   @Override

scio-core/src/main/java/com/spotify/scio/transforms/PipeDoFn.java

@@ -186,12 +186,12 @@ private Optional<RuntimeException> runCommands(String stage, List<String[]> comm
   }
 
   @StartBundle
-  public void startBundle() {
+  public void startBundle(StartBundleContext context) {
     isNewBundle = true;
   }
 
   @FinishBundle
-  public void finishBundle() {
+  public void finishBundle(FinishBundleContext context) {
     try {
       stdIn.close();
       int exitCode = pipeProcess.waitFor();

scio-test/src/test/scala/com/spotify/scio/transforms/AsyncDoFnSpec.scala

@@ -182,7 +182,7 @@ abstract class AsyncDoFnTester[P[_], F[_]] extends BaseDoFnTester {
   private val fn = {
     val f = newDoFn
     f.setup()
-    f.startBundle()
+    f.startBundle(null)
     f
   }
 
@@ -238,7 +238,7 @@ abstract class AsyncDoFnTester[P[_], F[_]] extends BaseDoFnTester {
     nextElement = 0
     pending.clear()
     outputBuffer.clear()
-    fn.startBundle()
+    fn.startBundle(null)
 
     result
   }