Merge pull request #8103 from boyuanzz/go_sdk

aaltay · web-flow · commit 4da6e93cb6fa · 2019-03-20T15:13:06.000-07:00
Reset Datasource counter to 0 in FinishBundle
diff --git a/sdks/go/pkg/beam/core/runtime/exec/datasource.go b/sdks/go/pkg/beam/core/runtime/exec/datasource.go
@@ -174,7 +174,9 @@ func (n *DataSource) Process(ctx context.Context) error {
 func (n *DataSource) FinishBundle(ctx context.Context) error {
 	log.Infof(ctx, "DataSource: %d elements in %d ns", atomic.LoadInt64(&n.count), time.Now().Sub(n.start))
 	n.source = nil
-	return n.Out.FinishBundle(ctx)
+	err := n.Out.FinishBundle(ctx)
+	atomic.StoreInt64(&n.count, 0)
+	return err
 }
 
 func (n *DataSource) Down(ctx context.Context) error {
