You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by dh...@apache.org on 2016/08/30 04:21:58 UTC

[2/2] incubator-beam git commit: take advantage of setup/teardown for KafkaWriter

take advantage of setup/teardown for KafkaWriter


Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/56c70027
Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/56c70027
Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/56c70027

Branch: refs/heads/master
Commit: 56c70027424cf665099ae64aca32f8076222033d
Parents: c948226
Author: Raghu Angadi <ra...@google.com>
Authored: Mon Aug 29 13:09:58 2016 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Mon Aug 29 21:21:46 2016 -0700

----------------------------------------------------------------------
 .../org/apache/beam/sdk/io/kafka/KafkaIO.java   | 25 ++++++++++----------
 1 file changed, 12 insertions(+), 13 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/56c70027/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index f16a5ab..6769b31 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -1325,16 +1325,12 @@ public class KafkaIO {
 
   private static class KafkaWriter<K, V> extends DoFn<KV<K, V>, Void> {
 
-    @StartBundle
-    public void startBundle(Context c) throws Exception {
-      // Producer initialization is fairly costly. Move this to future initialization api to avoid
-      // creating a producer for each bundle.
-      if (producer == null) {
-        if (producerFactoryFnOpt.isPresent()) {
-           producer = producerFactoryFnOpt.get().apply(producerConfig);
-        } else {
-          producer = new KafkaProducer<K, V>(producerConfig);
-        }
+    @Setup
+    public void setup() {
+      if (producerFactoryFnOpt.isPresent()) {
+        producer = producerFactoryFnOpt.get().apply(producerConfig);
+      } else {
+        producer = new KafkaProducer<K, V>(producerConfig);
       }
     }
 
@@ -1349,13 +1345,16 @@ public class KafkaIO {
     }
 
     @FinishBundle
-    public void finishBundle(Context c) throws Exception {
+    public void finishBundle(Context c) throws IOException {
       producer.flush();
-      producer.close();
-      producer = null;
       checkForFailures();
     }
 
+    @Teardown
+    public void teardown() {
+      producer.close();
+    }
+
     ///////////////////////////////////////////////////////////////////////////////////
 
     private final String topic;