[BEAM-14484] Improve behavior surrounding primary roots in self-checkpointing

sdks/go/pkg/beam/core/runtime/exec/datasource.go

@@ -29,6 +29,8 @@ import (
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/sdf"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/ioutilx"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/rtrackers/wrappedbounded"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/log"
 )
 
 // DataSource is a Root execution unit.
@@ -348,16 +350,25 @@ func (n *DataSource) makeEncodeElms() func([]*FullValue) ([][]byte, error) {
 	return encodeElms
 }
 
-func getRTrackerFromRoot(root *FullValue) (sdf.BoundableRTracker, float64, bool) {
-	tracker, ok := root.Elm.(*FullValue).Elm2.(*FullValue).Elm.(sdf.BoundableRTracker)
+func getBoundedRTrackerFromRoot(root *FullValue) (sdf.BoundableRTracker, float64, bool) {
+	tElm := root.Elm.(*FullValue).Elm2.(*FullValue).Elm
+	tracker, ok := tElm.(sdf.RTracker)
 	if !ok {
+		log.Warnf(context.Background(), "expected type sdf.RTracker, got type %T", tElm)
 		return nil, -1.0, false
 	}
+	boundTracker, ok := tracker.(sdf.BoundableRTracker)
+	if !ok {
+		log.Warn(context.Background(), "expected type sdf.BoundableRTracker; ensure that the RTracker implements IsBounded()")
+		// Assume an RTracker that does not implement IsBounded() will always be bounded, wrap so it can be used.
+		boundTracker = wrappedbounded.NewTracker(tracker)
+	}
 	size, ok := root.Elm2.(float64)
 	if !ok {
+		log.Warnf(context.Background(), "expected size to be type float64, got type %T", root.Elm2)
 		return nil, -1.0, false
 	}
-	return tracker, size, true
+	return boundTracker, size, true
 }
 
 // Checkpoint attempts to split an SDF that has self-checkpointed (e.g. returned a
@@ -385,11 +396,16 @@ func (n *DataSource) Checkpoint() (SplitResult, time.Duration, bool, error) {
 	if err != nil {
 		return SplitResult{}, -1 * time.Minute, false, err
 	}
+	if len(rs) == 0 {
+		return SplitResult{}, -1 * time.Minute, false, nil
+	}
 	if len(ps) != 0 {
+		// Expected structure of the root FullValue is KV<KV<Elm, KV<BoundedRTracker, watermarkEstimatorState>>, Size, (Timestamp?, Windows?)>
 		for _, root := range ps {
-			tracker, size, ok := getRTrackerFromRoot(root)
+			tracker, size, ok := getBoundedRTrackerFromRoot(root)
 			// If type assertion didn't return a BoundableRTracker, we move on.
 			if !ok {
+				log.Warnf(context.Background(), "got unexpected primary root contents %v, please check the output of the restriction tracker's TrySplit() function", root)
 				continue
 			}
 			if !tracker.IsBounded() || size > 0.00001 {

sdks/go/pkg/beam/core/runtime/exec/datasource_test.go

@@ -960,7 +960,7 @@ func TestGetRTrackerFromRoot(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			root := constructRootFullValue(test.inRt, test.inSize)
-			tracker, size, ok := getRTrackerFromRoot(root)
+			tracker, size, ok := getBoundedRTrackerFromRoot(root)
 
 			if test.valid {
 				if !ok {

sdks/go/pkg/beam/core/runtime/graphx/translate.go

@@ -43,6 +43,7 @@ const (
 	URNReshuffle     = "beam:transform:reshuffle:v1"
 	URNCombinePerKey = "beam:transform:combine_per_key:v1"
 	URNWindow        = "beam:transform:window_into:v1"
+	URNTruncate      = "beam:transform:sdf_truncate_sized_restrictions_v1"
 
 	URNIterableSideInput = "beam:side_input:iterable:v1"
 	URNMultimapSideInput = "beam:side_input:multimap:v1"

sdks/go/pkg/beam/io/rtrackers/wrappedbounded/wrappedbounded.go

@@ -0,0 +1,76 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package offsetrange defines a restriction and restriction tracker for offset
+// ranges. An offset range is just a range, with a start and end, that can
+// begin at an offset, and is commonly used to represent byte ranges for files
+// or indices for iterable containers.
+
+package wrappedbounded
+
+import "github.com/apache/beam/sdks/v2/go/pkg/beam/core/sdf"
+
+// Tracker wraps an implementation of an RTracker and adds an IsBounded() function
+// that returns true in order to allow RTrackers to be handled as bounded BoundableRTrackers
+// if necessary (like in self-checkpointing evaluation.)
+type Tracker struct {
+	baseTracker sdf.RTracker
+}
+
+// TryClaim attempts to claim a block of work from the underlying RTracker's restriction.
+func (t *Tracker) TryClaim(pos interface{}) (ok bool) {
+	return t.baseTracker.TryClaim(pos)
+}
+
+// GetError returns an error from the underlying RTracker if it has stopped executing. Returns nil
+// if none has occurred.
+func (t *Tracker) GetError() error {
+	return t.baseTracker.GetError()
+}
+
+// TrySplit splits the underlying RTracker's restriction into a primary (work that is currently executing)
+// and a residual (work that will be split off and resumed later.)
+func (t *Tracker) TrySplit(fraction float64) (primary, residual interface{}, err error) {
+	return t.baseTracker.TrySplit(fraction)
+}
+
+// GetProgress returns two abstract scalars representing the amount of work done and the remaining work
+// left in the underlying RTracker. These are unitless values, only used to estimate work in relation to
+// each other.
+func (t *Tracker) GetProgress() (done float64, remaining float64) {
+	return t.baseTracker.GetProgress()
+}
+
+// IsDone() returns a boolean indicating if the work represented by the underlying RTracker has
+// been completed.
+func (t *Tracker) IsDone() bool {
+	return t.baseTracker.IsDone()
+}
+
+// GetRestriction returns the restriction maintained by the underlying RTracker.
+func (t *Tracker) GetRestriction() interface{} {
+	return t.baseTracker.GetRestriction()
+}
+
+// IsBounded returns true, indicating that the underlying RTracker represents a bounded
+// amount of work.
+func (t *Tracker) IsBounded() bool {
+	return true
+}
+
+// NewTracker is a constructor for an RTracker that wraps another RTracker into a BoundedRTracker.
+func NewTracker(underlying sdf.RTracker) *Tracker {
+	return &Tracker{baseTracker: underlying}
+}

sdks/go/test/integration/primitives/checkpointing.go

@@ -53,7 +53,7 @@ func (fn *selfCheckpointingDoFn) RestrictionSize(_ []byte, rest offsetrange.Rest
 // SplitRestriction modifies the offsetrange.Restriction's sized restriction function to produce a size-zero restriction
 // at the end of execution.
 func (fn *selfCheckpointingDoFn) SplitRestriction(_ []byte, rest offsetrange.Restriction) []offsetrange.Restriction {
-	size := int64(1)
+	size := int64(10)
 	s := rest.Start
 	var splits []offsetrange.Restriction
 	for e := s + size; e <= rest.End; s, e = e, e+size {
@@ -68,19 +68,24 @@ func (fn *selfCheckpointingDoFn) SplitRestriction(_ []byte, rest offsetrange.Res
 func (fn *selfCheckpointingDoFn) ProcessElement(rt *sdf.LockRTracker, _ []byte, emit func(int64)) sdf.ProcessContinuation {
 	position := rt.GetRestriction().(offsetrange.Restriction).Start
 
+	counter := 0
 	for {
 		if rt.TryClaim(position) {
 			// Successful claim, emit the value and move on.
 			emit(position)
 			position++
-			return sdf.ResumeProcessingIn(1 * time.Second)
+			counter++
 		} else if rt.GetError() != nil || rt.IsDone() {
 			// Stop processing on error or completion
 			return sdf.StopProcessing()
 		} else {
-			// Failed to claim but no error, resume later.
+			// Resume later.
 			return sdf.ResumeProcessingIn(5 * time.Second)
 		}
+
+		if counter >= 10 {
+			return sdf.ResumeProcessingIn(1 * time.Second)
+		}
 	}
 }