Fix bug with sdf bundle finalization

apache · Jun 2, 2022 · 04b24dc · 04b24dc
1 parent ae52acb
commit 04b24dc
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 3 deletions.
diff --git a/sdks/go/examples/snippets/04transforms.go b/sdks/go/examples/snippets/04transforms.go
@@ -24,6 +24,7 @@ import (
 	"time"
 
 	"github.com/apache/beam/sdks/v2/go/pkg/beam"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/sdf"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/rtrackers/offsetrange"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/transforms/stats"
@@ -96,7 +97,7 @@ type weDoFn struct{}
 
 // [START bundlefinalization_simplecallback]
 
-func (fn *splittableDoFn) ProcessElement(bf beam.BundleFinalization, element string) {
+func (fn *splittableDoFn) ProcessElement(bf beam.BundleFinalization, rt *sdf.LockRTracker, element string) {
 	// ... produce output ...
 
 	bf.RegisterCallback(5*time.Minute, func() error {

diff --git a/sdks/go/pkg/beam/core/runtime/exec/fn.go b/sdks/go/pkg/beam/core/runtime/exec/fn.go
@@ -52,6 +52,10 @@ type bundleFinalizer struct {
 	lastValidCallback time.Time // Used to track when we can safely gc the bundleFinalizer
 }
 
+type needsBundleFinalization interface {
+	AttachFinalizer(*bundleFinalizer)
+}
+
 // RegisterCallback is used to register callbacks during DoFn execution.
 func (bf *bundleFinalizer) RegisterCallback(t time.Duration, cb func() error) {
 	callback := bundleFinalizationCallback{

diff --git a/sdks/go/pkg/beam/core/runtime/exec/pardo.go b/sdks/go/pkg/beam/core/runtime/exec/pardo.go
@@ -99,6 +99,10 @@ func (n *ParDo) Up(ctx context.Context) error {
 	return nil
 }
 
+func (n *ParDo) AttachFinalizer(bf *bundleFinalizer) {
+	n.bf = bf
+}
+
 // StartBundle does pre-bundle processing operation for the DoFn.
 func (n *ParDo) StartBundle(ctx context.Context, id string, data DataContext) error {
 	if n.status != Up {

diff --git a/sdks/go/pkg/beam/core/runtime/exec/plan.go b/sdks/go/pkg/beam/core/runtime/exec/plan.go
@@ -66,8 +66,8 @@ func NewPlan(id string, units []Unit) (*Plan, error) {
 		if p, ok := u.(*PCollection); ok {
 			pcols = append(pcols, p)
 		}
-		if p, ok := u.(*ParDo); ok {
-			p.bf = &bf
+		if p, ok := u.(needsBundleFinalization); ok {
+			p.AttachFinalizer(&bf)
 		}
 	}
 	if len(roots) == 0 {

diff --git a/sdks/go/pkg/beam/core/runtime/exec/sdf.go b/sdks/go/pkg/beam/core/runtime/exec/sdf.go
@@ -437,6 +437,10 @@ func (n *ProcessSizedElementsAndRestrictions) Up(ctx context.Context) error {
 	return n.PDo.Up(ctx)
 }
 
+func (n *ProcessSizedElementsAndRestrictions) AttachFinalizer(bf *bundleFinalizer) {
+	n.PDo.bf = bf
+}
+
 // StartBundle calls the ParDo's StartBundle method.
 func (n *ProcessSizedElementsAndRestrictions) StartBundle(ctx context.Context, id string, data DataContext) error {
 	return n.PDo.StartBundle(ctx, id, data)
@@ -950,6 +954,10 @@ func (n *SdfFallback) Up(ctx context.Context) error {
 	return n.PDo.Up(ctx)
 }
 
+func (n *SdfFallback) AttachFinalizer(bf *bundleFinalizer) {
+	n.PDo.bf = bf
+}
+
 // StartBundle calls the ParDo's StartBundle method.
 func (n *SdfFallback) StartBundle(ctx context.Context, id string, data DataContext) error {
 	return n.PDo.StartBundle(ctx, id, data)