You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by GitBox <gi...@apache.org> on 2020/05/20 19:09:10 UTC

[GitHub] [beam] youngoli opened a new pull request #11763: [BEAM-9978] Adding functionality and tests to Go offset range tracker.

youngoli opened a new pull request #11763:
URL: https://github.com/apache/beam/pull/11763


   Pretty simple. Moves some commonly desired behaviors out of the SDF
   code and into the offset range tracker/restriction code + adds tests.
   
   ------------------------
   
   Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
   
    - [x] [**Choose reviewer(s)**](https://beam.apache.org/contribute/#make-your-change) and mention them in a comment (`R: @username`).
    - [x] Format the pull request title like `[BEAM-XXX] Fixes bug in ApproximateQuantiles`, where you replace `BEAM-XXX` with the appropriate JIRA issue, if applicable. This will automatically link the pull request to the issue.
    - [ ] Update `CHANGES.md` with noteworthy changes.
    - [x] If this contribution is large, please file an Apache [Individual Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
   
   See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://beam.apache.org/contribute/#make-reviewers-job-easier).
   
   Post-Commit Tests Status (on master branch)
   ------------------------------------------------------------------------------------------------
   
   Lang | SDK | Apex | Dataflow | Flink | Gearpump | Samza | Spark
   --- | --- | --- | --- | --- | --- | --- | ---
   Go | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Go/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Go/lastCompletedBuild/) | --- | --- | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Go_VR_Flink/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Go_VR_Flink/lastCompletedBuild/) | --- | --- | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Go_VR_Spark/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Go_VR_Spark/lastCompletedBuild/)
   Java | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Java/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java/lastCompletedBuild/) | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Apex/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Apex/lastCompletedBuild/) | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow_Java11/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow_Java11/lastCompletedBuild/) | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink_Java11/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink_Java11/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Flink_Batch/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Flink_Batch/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Flink_Streaming/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Flink_Streaming/lastCompletedBuild/) | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Gearpump/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Gearpump/lastCompletedBuild/) | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Samza/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Samza/lastCompletedBuild/) | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Spark/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Spark/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Spark_Batch/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Spark_Batch/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming/lastCompletedBuild/)
   Python | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Python2/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Python2/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Python35/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Python35/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Python36/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Python36/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Python37/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Python37/lastCompletedBuild/) | --- | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Py_VR_Dataflow/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Py_VR_Dataflow/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Py_VR_Dataflow_V2/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Py_VR_Dataflow_V2/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Py_ValCont/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Py_ValCont/lastCompletedBuild/) | [![Build Status](https://builds.apache.org/job/beam_PreCommit_Python2_PVR_Flink_Cron/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PreCommit_Python2_PVR_Flink_Cron/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PostCommit_Python35_VR_Flink/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Python35_VR_Flink/lastCompletedBuild/) | --- | --- | [![Build Status](https://builds.apache.org/job/beam_PostCommit_Python_VR_Spark/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_Python_VR_Spark/lastCompletedBuild/)
   XLang | --- | --- | --- | [![Build Status](https://builds.apache.org/job/beam_PostCommit_XVR_Flink/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_XVR_Flink/lastCompletedBuild/) | --- | --- | [![Build Status](https://builds.apache.org/job/beam_PostCommit_XVR_Spark/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PostCommit_XVR_Spark/lastCompletedBuild/)
   
   Pre-Commit Tests Status (on master branch)
   ------------------------------------------------------------------------------------------------
   
   --- |Java | Python | Go | Website
   --- | --- | --- | --- | ---
   Non-portable | [![Build Status](https://builds.apache.org/job/beam_PreCommit_Java_Cron/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PreCommit_Java_Cron/lastCompletedBuild/) | [![Build Status](https://builds.apache.org/job/beam_PreCommit_Python_Cron/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PreCommit_Python_Cron/lastCompletedBuild/)<br>[![Build Status](https://builds.apache.org/job/beam_PreCommit_PythonLint_Cron/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PreCommit_PythonLint_Cron/lastCompletedBuild/) | [![Build Status](https://builds.apache.org/job/beam_PreCommit_Go_Cron/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PreCommit_Go_Cron/lastCompletedBuild/) | [![Build Status](https://builds.apache.org/job/beam_PreCommit_Website_Cron/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PreCommit_Website_Cron/lastCompletedBuild/) 
   Portable | --- | [![Build Status](https://builds.apache.org/job/beam_PreCommit_Portable_Python_Cron/lastCompletedBuild/badge/icon)](https://builds.apache.org/job/beam_PreCommit_Portable_Python_Cron/lastCompletedBuild/) | --- | ---
   
   See [.test-infra/jenkins/README](https://github.com/apache/beam/blob/master/.test-infra/jenkins/README.md) for trigger phrase, status and link of all Jenkins jobs.
   


----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [beam] youngoli merged pull request #11763: [BEAM-9978] Adding functionality and tests to Go offset range tracker.

Posted by GitBox <gi...@apache.org>.
youngoli merged pull request #11763:
URL: https://github.com/apache/beam/pull/11763


   


----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [beam] youngoli commented on a change in pull request #11763: [BEAM-9978] Adding functionality and tests to Go offset range tracker.

Posted by GitBox <gi...@apache.org>.
youngoli commented on a change in pull request #11763:
URL: https://github.com/apache/beam/pull/11763#discussion_r429026079



##########
File path: sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange_test.go
##########
@@ -0,0 +1,212 @@
+// 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
+
+import (
+	"fmt"
+	"github.com/google/go-cmp/cmp"
+	"testing"
+)
+
+// TestRestriction_EvenSplits tests various splits and checks that they all
+// follow the contract for EvenSplits. This means that all restrictions are
+// evenly split, that each restriction has at least one element, and that each
+// element is present in the split restrictions.
+func TestRestriction_EvenSplits(t *testing.T) {
+	tests := []struct {
+		rest Restriction
+		num  int64
+	}{
+		{rest: Restriction{Start: 0, End: 21}, num: 4},
+		{rest: Restriction{Start: 21, End: 42}, num: 4},
+		{rest: Restriction{Start: 0, End: 5}, num: 10},
+		{rest: Restriction{Start: 0, End: 21}, num: -1},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(fmt.Sprintf("(rest[%v, %v], splits = %v)",
+			test.rest.Start, test.rest.End, test.num), func(t *testing.T) {
+			r := test.rest
+
+			// Get the minimum size that a split restriction can be. Max size
+			// should be min + 1. This way we can check the size of each split.
+			num := test.num
+			if num <= 1 {
+				num = 1
+			}
+			min := (r.End - r.Start) / num
+
+			splits := r.EvenSplits(test.num)
+			prevEnd := r.Start
+			for _, split := range splits {
+				size := split.End - split.Start
+				// Check: Each restriction has at least 1 element.
+				if size == 0 {
+					t.Errorf("split restriction [%v, %v] is empty, size must be greater than 0.",
+						split.Start, split.End)
+				}
+				// Check: Restrictions are evenly split.
+				if size != min && size != min+1 {
+					t.Errorf("split restriction [%v, %v] has unexpected size. got: %v, want: %v or %v",
+						split.Start, split.End, size, min, min+1)
+				}
+				// Check: All elements are still in a split restrictions. This
+				// logic assumes that the splits are returned in order which
+				// isn't guaranteed by EvenSplits, but this check is way easier
+				// with the assumption.
+				if split.Start != prevEnd {
+					t.Errorf("restriction range [%v, %v] missing after splits.",
+						prevEnd, split.Start)
+				} else {
+					prevEnd = split.End
+				}
+			}
+			if prevEnd != r.End {
+				t.Errorf("restriction range [%v, %v] missing after splits.",
+					prevEnd, r.End)
+			}
+		})
+	}
+}
+
+// TestTracker_TryClaim validates both success and failure cases for TryClaim.
+func TestTracker_TryClaim(t *testing.T) {
+	// Test that TryClaim works as expected when called correctly.
+	t.Run("Correctness", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 1, 2, 3}},
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{15, 20, 50}},
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)
+					// If TryClaim returns false, check if there was an error.
+					if !ok && !rt.IsDone() {
+						t.Fatalf("tracker claiming %v failed, error: %v", pos, rt.GetError())
+					}
+				}
+			})
+		}
+	})
+
+	// Test that each invalid error case actually results in an error.
+	t.Run("Errors", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			// Claiming backwards.
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 2, 1}},
+			// Claiming before start of restriction.
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{8}},
+			// Claiming after tracker signalled to stop.
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4, 5}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)
+					// Finish successfully if we got an error.
+					if !ok && !rt.IsDone() && rt.GetError() != nil {
+						return
+					}
+				}
+				t.Fatal("tracker did not fail on invalid claim")
+			})
+		}
+	})
+}
+
+// TestTracker_TrySplit tests that TrySplit follows its contract, meaning that
+// splits don't lose any elements, split fractions are clamped to 0 or 1, and
+// that TrySplit always splits at the nearest integer greater than the given
+// fraction.
+func TestTracker_TrySplit(t *testing.T) {
+	tests := []struct {
+		rest     Restriction
+		claimed  int64
+		fraction float64
+		splitPt  int64 // Index where we want the split to happen.

Review comment:
       The range is actually half-closed range, so that split is the "first element that shouldn't be processed" of the primary, and the "first element to process" in the residual. But agreed, specified in the comment.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [beam] youngoli commented on pull request #11763: [BEAM-9978] Adding functionality and tests to Go offset range tracker.

Posted by GitBox <gi...@apache.org>.
youngoli commented on pull request #11763:
URL: https://github.com/apache/beam/pull/11763#issuecomment-631668755


   R: @lostluck 


----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [beam] youngoli commented on a change in pull request #11763: [BEAM-9978] Adding functionality and tests to Go offset range tracker.

Posted by GitBox <gi...@apache.org>.
youngoli commented on a change in pull request #11763:
URL: https://github.com/apache/beam/pull/11763#discussion_r429025895



##########
File path: sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange_test.go
##########
@@ -0,0 +1,212 @@
+// 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
+
+import (
+	"fmt"
+	"github.com/google/go-cmp/cmp"
+	"testing"
+)
+
+// TestRestriction_EvenSplits tests various splits and checks that they all
+// follow the contract for EvenSplits. This means that all restrictions are
+// evenly split, that each restriction has at least one element, and that each
+// element is present in the split restrictions.
+func TestRestriction_EvenSplits(t *testing.T) {
+	tests := []struct {
+		rest Restriction
+		num  int64
+	}{
+		{rest: Restriction{Start: 0, End: 21}, num: 4},
+		{rest: Restriction{Start: 21, End: 42}, num: 4},
+		{rest: Restriction{Start: 0, End: 5}, num: 10},
+		{rest: Restriction{Start: 0, End: 21}, num: -1},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(fmt.Sprintf("(rest[%v, %v], splits = %v)",
+			test.rest.Start, test.rest.End, test.num), func(t *testing.T) {
+			r := test.rest
+
+			// Get the minimum size that a split restriction can be. Max size
+			// should be min + 1. This way we can check the size of each split.
+			num := test.num
+			if num <= 1 {
+				num = 1
+			}
+			min := (r.End - r.Start) / num
+
+			splits := r.EvenSplits(test.num)
+			prevEnd := r.Start
+			for _, split := range splits {
+				size := split.End - split.Start
+				// Check: Each restriction has at least 1 element.
+				if size == 0 {
+					t.Errorf("split restriction [%v, %v] is empty, size must be greater than 0.",
+						split.Start, split.End)
+				}
+				// Check: Restrictions are evenly split.
+				if size != min && size != min+1 {
+					t.Errorf("split restriction [%v, %v] has unexpected size. got: %v, want: %v or %v",
+						split.Start, split.End, size, min, min+1)
+				}
+				// Check: All elements are still in a split restrictions. This
+				// logic assumes that the splits are returned in order which
+				// isn't guaranteed by EvenSplits, but this check is way easier
+				// with the assumption.
+				if split.Start != prevEnd {
+					t.Errorf("restriction range [%v, %v] missing after splits.",
+						prevEnd, split.Start)
+				} else {
+					prevEnd = split.End
+				}
+			}
+			if prevEnd != r.End {
+				t.Errorf("restriction range [%v, %v] missing after splits.",
+					prevEnd, r.End)
+			}
+		})
+	}
+}
+
+// TestTracker_TryClaim validates both success and failure cases for TryClaim.
+func TestTracker_TryClaim(t *testing.T) {
+	// Test that TryClaim works as expected when called correctly.
+	t.Run("Correctness", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 1, 2, 3}},
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{15, 20, 50}},
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)

Review comment:
       Done.

##########
File path: sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange_test.go
##########
@@ -0,0 +1,212 @@
+// 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
+
+import (
+	"fmt"
+	"github.com/google/go-cmp/cmp"
+	"testing"
+)
+
+// TestRestriction_EvenSplits tests various splits and checks that they all
+// follow the contract for EvenSplits. This means that all restrictions are
+// evenly split, that each restriction has at least one element, and that each
+// element is present in the split restrictions.
+func TestRestriction_EvenSplits(t *testing.T) {
+	tests := []struct {
+		rest Restriction
+		num  int64
+	}{
+		{rest: Restriction{Start: 0, End: 21}, num: 4},
+		{rest: Restriction{Start: 21, End: 42}, num: 4},
+		{rest: Restriction{Start: 0, End: 5}, num: 10},
+		{rest: Restriction{Start: 0, End: 21}, num: -1},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(fmt.Sprintf("(rest[%v, %v], splits = %v)",
+			test.rest.Start, test.rest.End, test.num), func(t *testing.T) {
+			r := test.rest
+
+			// Get the minimum size that a split restriction can be. Max size
+			// should be min + 1. This way we can check the size of each split.
+			num := test.num
+			if num <= 1 {
+				num = 1
+			}
+			min := (r.End - r.Start) / num
+
+			splits := r.EvenSplits(test.num)
+			prevEnd := r.Start
+			for _, split := range splits {
+				size := split.End - split.Start
+				// Check: Each restriction has at least 1 element.
+				if size == 0 {
+					t.Errorf("split restriction [%v, %v] is empty, size must be greater than 0.",
+						split.Start, split.End)
+				}
+				// Check: Restrictions are evenly split.
+				if size != min && size != min+1 {
+					t.Errorf("split restriction [%v, %v] has unexpected size. got: %v, want: %v or %v",
+						split.Start, split.End, size, min, min+1)
+				}
+				// Check: All elements are still in a split restrictions. This
+				// logic assumes that the splits are returned in order which
+				// isn't guaranteed by EvenSplits, but this check is way easier
+				// with the assumption.
+				if split.Start != prevEnd {
+					t.Errorf("restriction range [%v, %v] missing after splits.",
+						prevEnd, split.Start)
+				} else {
+					prevEnd = split.End
+				}
+			}
+			if prevEnd != r.End {
+				t.Errorf("restriction range [%v, %v] missing after splits.",
+					prevEnd, r.End)
+			}
+		})
+	}
+}
+
+// TestTracker_TryClaim validates both success and failure cases for TryClaim.
+func TestTracker_TryClaim(t *testing.T) {
+	// Test that TryClaim works as expected when called correctly.
+	t.Run("Correctness", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 1, 2, 3}},
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{15, 20, 50}},
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)
+					// If TryClaim returns false, check if there was an error.
+					if !ok && !rt.IsDone() {
+						t.Fatalf("tracker claiming %v failed, error: %v", pos, rt.GetError())
+					}
+				}
+			})
+		}
+	})
+
+	// Test that each invalid error case actually results in an error.
+	t.Run("Errors", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			// Claiming backwards.
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 2, 1}},
+			// Claiming before start of restriction.
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{8}},
+			// Claiming after tracker signalled to stop.
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4, 5}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)
+					// Finish successfully if we got an error.
+					if !ok && !rt.IsDone() && rt.GetError() != nil {
+						return
+					}
+				}
+				t.Fatal("tracker did not fail on invalid claim")
+			})
+		}
+	})
+}
+
+// TestTracker_TrySplit tests that TrySplit follows its contract, meaning that
+// splits don't lose any elements, split fractions are clamped to 0 or 1, and
+// that TrySplit always splits at the nearest integer greater than the given
+// fraction.
+func TestTracker_TrySplit(t *testing.T) {
+	tests := []struct {
+		rest     Restriction
+		claimed  int64
+		fraction float64
+		splitPt  int64 // Index where we want the split to happen.
+	}{
+		{
+			rest:     Restriction{Start: 0, End: 1},
+			claimed:  0,
+			fraction: 0.5,
+			splitPt:  1,
+		},
+		{
+			rest:     Restriction{Start: 0, End: 5},
+			claimed:  0,
+			fraction: 0.5,
+			splitPt:  3,
+		},
+		{
+			rest:     Restriction{Start: 0, End: 10},
+			claimed:  5,
+			fraction: 0.5,
+			splitPt:  8,
+		},
+		{
+			rest:     Restriction{Start: 0, End: 10},
+			claimed:  5,
+			fraction: -0.5,
+			splitPt:  5,
+		},
+		{
+			rest:     Restriction{Start: 0, End: 10},
+			claimed:  5,
+			fraction: 1.5,
+			splitPt:  10,
+		},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(fmt.Sprintf("(split at %v of [%v, %v])",
+			test.fraction, test.claimed, test.rest.End), func(t *testing.T) {
+			rt := NewTracker(test.rest)
+			ok := rt.TryClaim(test.claimed)
+			if !ok {

Review comment:
       Done.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [beam] lostluck commented on a change in pull request #11763: [BEAM-9978] Adding functionality and tests to Go offset range tracker.

Posted by GitBox <gi...@apache.org>.
lostluck commented on a change in pull request #11763:
URL: https://github.com/apache/beam/pull/11763#discussion_r428305778



##########
File path: sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange_test.go
##########
@@ -0,0 +1,212 @@
+// 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
+
+import (
+	"fmt"
+	"github.com/google/go-cmp/cmp"
+	"testing"
+)
+
+// TestRestriction_EvenSplits tests various splits and checks that they all
+// follow the contract for EvenSplits. This means that all restrictions are
+// evenly split, that each restriction has at least one element, and that each
+// element is present in the split restrictions.
+func TestRestriction_EvenSplits(t *testing.T) {
+	tests := []struct {
+		rest Restriction
+		num  int64
+	}{
+		{rest: Restriction{Start: 0, End: 21}, num: 4},
+		{rest: Restriction{Start: 21, End: 42}, num: 4},
+		{rest: Restriction{Start: 0, End: 5}, num: 10},
+		{rest: Restriction{Start: 0, End: 21}, num: -1},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(fmt.Sprintf("(rest[%v, %v], splits = %v)",
+			test.rest.Start, test.rest.End, test.num), func(t *testing.T) {
+			r := test.rest
+
+			// Get the minimum size that a split restriction can be. Max size
+			// should be min + 1. This way we can check the size of each split.
+			num := test.num
+			if num <= 1 {
+				num = 1
+			}
+			min := (r.End - r.Start) / num
+
+			splits := r.EvenSplits(test.num)
+			prevEnd := r.Start
+			for _, split := range splits {
+				size := split.End - split.Start
+				// Check: Each restriction has at least 1 element.
+				if size == 0 {
+					t.Errorf("split restriction [%v, %v] is empty, size must be greater than 0.",
+						split.Start, split.End)
+				}
+				// Check: Restrictions are evenly split.
+				if size != min && size != min+1 {
+					t.Errorf("split restriction [%v, %v] has unexpected size. got: %v, want: %v or %v",
+						split.Start, split.End, size, min, min+1)
+				}
+				// Check: All elements are still in a split restrictions. This
+				// logic assumes that the splits are returned in order which
+				// isn't guaranteed by EvenSplits, but this check is way easier
+				// with the assumption.
+				if split.Start != prevEnd {
+					t.Errorf("restriction range [%v, %v] missing after splits.",
+						prevEnd, split.Start)
+				} else {
+					prevEnd = split.End
+				}
+			}
+			if prevEnd != r.End {
+				t.Errorf("restriction range [%v, %v] missing after splits.",
+					prevEnd, r.End)
+			}
+		})
+	}
+}
+
+// TestTracker_TryClaim validates both success and failure cases for TryClaim.
+func TestTracker_TryClaim(t *testing.T) {
+	// Test that TryClaim works as expected when called correctly.
+	t.Run("Correctness", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 1, 2, 3}},
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{15, 20, 50}},
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)

Review comment:
       Stylistically, I'd move this inline into the if instread of putting into the variable.

##########
File path: sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange_test.go
##########
@@ -0,0 +1,212 @@
+// 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
+
+import (
+	"fmt"
+	"github.com/google/go-cmp/cmp"
+	"testing"
+)
+
+// TestRestriction_EvenSplits tests various splits and checks that they all
+// follow the contract for EvenSplits. This means that all restrictions are
+// evenly split, that each restriction has at least one element, and that each
+// element is present in the split restrictions.
+func TestRestriction_EvenSplits(t *testing.T) {
+	tests := []struct {
+		rest Restriction
+		num  int64
+	}{
+		{rest: Restriction{Start: 0, End: 21}, num: 4},
+		{rest: Restriction{Start: 21, End: 42}, num: 4},
+		{rest: Restriction{Start: 0, End: 5}, num: 10},
+		{rest: Restriction{Start: 0, End: 21}, num: -1},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(fmt.Sprintf("(rest[%v, %v], splits = %v)",
+			test.rest.Start, test.rest.End, test.num), func(t *testing.T) {
+			r := test.rest
+
+			// Get the minimum size that a split restriction can be. Max size
+			// should be min + 1. This way we can check the size of each split.
+			num := test.num
+			if num <= 1 {
+				num = 1
+			}
+			min := (r.End - r.Start) / num
+
+			splits := r.EvenSplits(test.num)
+			prevEnd := r.Start
+			for _, split := range splits {
+				size := split.End - split.Start
+				// Check: Each restriction has at least 1 element.
+				if size == 0 {
+					t.Errorf("split restriction [%v, %v] is empty, size must be greater than 0.",
+						split.Start, split.End)
+				}
+				// Check: Restrictions are evenly split.
+				if size != min && size != min+1 {
+					t.Errorf("split restriction [%v, %v] has unexpected size. got: %v, want: %v or %v",
+						split.Start, split.End, size, min, min+1)
+				}
+				// Check: All elements are still in a split restrictions. This
+				// logic assumes that the splits are returned in order which
+				// isn't guaranteed by EvenSplits, but this check is way easier
+				// with the assumption.
+				if split.Start != prevEnd {
+					t.Errorf("restriction range [%v, %v] missing after splits.",
+						prevEnd, split.Start)
+				} else {
+					prevEnd = split.End
+				}
+			}
+			if prevEnd != r.End {
+				t.Errorf("restriction range [%v, %v] missing after splits.",
+					prevEnd, r.End)
+			}
+		})
+	}
+}
+
+// TestTracker_TryClaim validates both success and failure cases for TryClaim.
+func TestTracker_TryClaim(t *testing.T) {
+	// Test that TryClaim works as expected when called correctly.
+	t.Run("Correctness", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 1, 2, 3}},
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{15, 20, 50}},
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)
+					// If TryClaim returns false, check if there was an error.
+					if !ok && !rt.IsDone() {
+						t.Fatalf("tracker claiming %v failed, error: %v", pos, rt.GetError())
+					}
+				}
+			})
+		}
+	})
+
+	// Test that each invalid error case actually results in an error.
+	t.Run("Errors", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			// Claiming backwards.
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 2, 1}},
+			// Claiming before start of restriction.
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{8}},
+			// Claiming after tracker signalled to stop.
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4, 5}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)
+					// Finish successfully if we got an error.
+					if !ok && !rt.IsDone() && rt.GetError() != nil {
+						return
+					}
+				}
+				t.Fatal("tracker did not fail on invalid claim")
+			})
+		}
+	})
+}
+
+// TestTracker_TrySplit tests that TrySplit follows its contract, meaning that
+// splits don't lose any elements, split fractions are clamped to 0 or 1, and
+// that TrySplit always splits at the nearest integer greater than the given
+// fraction.
+func TestTracker_TrySplit(t *testing.T) {
+	tests := []struct {
+		rest     Restriction
+		claimed  int64
+		fraction float64
+		splitPt  int64 // Index where we want the split to happen.
+	}{
+		{
+			rest:     Restriction{Start: 0, End: 1},
+			claimed:  0,
+			fraction: 0.5,
+			splitPt:  1,
+		},
+		{
+			rest:     Restriction{Start: 0, End: 5},
+			claimed:  0,
+			fraction: 0.5,
+			splitPt:  3,
+		},
+		{
+			rest:     Restriction{Start: 0, End: 10},
+			claimed:  5,
+			fraction: 0.5,
+			splitPt:  8,
+		},
+		{
+			rest:     Restriction{Start: 0, End: 10},
+			claimed:  5,
+			fraction: -0.5,
+			splitPt:  5,
+		},
+		{
+			rest:     Restriction{Start: 0, End: 10},
+			claimed:  5,
+			fraction: 1.5,
+			splitPt:  10,
+		},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(fmt.Sprintf("(split at %v of [%v, %v])",
+			test.fraction, test.claimed, test.rest.End), func(t *testing.T) {
+			rt := NewTracker(test.rest)
+			ok := rt.TryClaim(test.claimed)
+			if !ok {

Review comment:
       Same comment here.

##########
File path: sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange_test.go
##########
@@ -0,0 +1,212 @@
+// 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
+
+import (
+	"fmt"
+	"github.com/google/go-cmp/cmp"
+	"testing"
+)
+
+// TestRestriction_EvenSplits tests various splits and checks that they all
+// follow the contract for EvenSplits. This means that all restrictions are
+// evenly split, that each restriction has at least one element, and that each
+// element is present in the split restrictions.
+func TestRestriction_EvenSplits(t *testing.T) {
+	tests := []struct {
+		rest Restriction
+		num  int64
+	}{
+		{rest: Restriction{Start: 0, End: 21}, num: 4},
+		{rest: Restriction{Start: 21, End: 42}, num: 4},
+		{rest: Restriction{Start: 0, End: 5}, num: 10},
+		{rest: Restriction{Start: 0, End: 21}, num: -1},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(fmt.Sprintf("(rest[%v, %v], splits = %v)",
+			test.rest.Start, test.rest.End, test.num), func(t *testing.T) {
+			r := test.rest
+
+			// Get the minimum size that a split restriction can be. Max size
+			// should be min + 1. This way we can check the size of each split.
+			num := test.num
+			if num <= 1 {
+				num = 1
+			}
+			min := (r.End - r.Start) / num
+
+			splits := r.EvenSplits(test.num)
+			prevEnd := r.Start
+			for _, split := range splits {
+				size := split.End - split.Start
+				// Check: Each restriction has at least 1 element.
+				if size == 0 {
+					t.Errorf("split restriction [%v, %v] is empty, size must be greater than 0.",
+						split.Start, split.End)
+				}
+				// Check: Restrictions are evenly split.
+				if size != min && size != min+1 {
+					t.Errorf("split restriction [%v, %v] has unexpected size. got: %v, want: %v or %v",
+						split.Start, split.End, size, min, min+1)
+				}
+				// Check: All elements are still in a split restrictions. This
+				// logic assumes that the splits are returned in order which
+				// isn't guaranteed by EvenSplits, but this check is way easier
+				// with the assumption.
+				if split.Start != prevEnd {
+					t.Errorf("restriction range [%v, %v] missing after splits.",
+						prevEnd, split.Start)
+				} else {
+					prevEnd = split.End
+				}
+			}
+			if prevEnd != r.End {
+				t.Errorf("restriction range [%v, %v] missing after splits.",
+					prevEnd, r.End)
+			}
+		})
+	}
+}
+
+// TestTracker_TryClaim validates both success and failure cases for TryClaim.
+func TestTracker_TryClaim(t *testing.T) {
+	// Test that TryClaim works as expected when called correctly.
+	t.Run("Correctness", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 1, 2, 3}},
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{15, 20, 50}},
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)
+					// If TryClaim returns false, check if there was an error.
+					if !ok && !rt.IsDone() {
+						t.Fatalf("tracker claiming %v failed, error: %v", pos, rt.GetError())
+					}
+				}
+			})
+		}
+	})
+
+	// Test that each invalid error case actually results in an error.
+	t.Run("Errors", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			// Claiming backwards.
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 2, 1}},
+			// Claiming before start of restriction.
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{8}},
+			// Claiming after tracker signalled to stop.
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4, 5}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)
+					// Finish successfully if we got an error.
+					if !ok && !rt.IsDone() && rt.GetError() != nil {
+						return
+					}
+				}
+				t.Fatal("tracker did not fail on invalid claim")
+			})
+		}
+	})
+}
+
+// TestTracker_TrySplit tests that TrySplit follows its contract, meaning that
+// splits don't lose any elements, split fractions are clamped to 0 or 1, and
+// that TrySplit always splits at the nearest integer greater than the given
+// fraction.
+func TestTracker_TrySplit(t *testing.T) {
+	tests := []struct {
+		rest     Restriction
+		claimed  int64
+		fraction float64
+		splitPt  int64 // Index where we want the split to happen.

Review comment:
       Is the splitPt here the "last element that should be processed" or "the first element that shouldn't be processed?"
   Said another way, is the index represented by splitPt part of the Primary or part of the Residual?

##########
File path: sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange_test.go
##########
@@ -0,0 +1,212 @@
+// 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
+
+import (
+	"fmt"
+	"github.com/google/go-cmp/cmp"
+	"testing"
+)
+
+// TestRestriction_EvenSplits tests various splits and checks that they all
+// follow the contract for EvenSplits. This means that all restrictions are
+// evenly split, that each restriction has at least one element, and that each
+// element is present in the split restrictions.
+func TestRestriction_EvenSplits(t *testing.T) {
+	tests := []struct {
+		rest Restriction
+		num  int64
+	}{
+		{rest: Restriction{Start: 0, End: 21}, num: 4},
+		{rest: Restriction{Start: 21, End: 42}, num: 4},
+		{rest: Restriction{Start: 0, End: 5}, num: 10},
+		{rest: Restriction{Start: 0, End: 21}, num: -1},
+	}
+	for _, test := range tests {
+		test := test
+		t.Run(fmt.Sprintf("(rest[%v, %v], splits = %v)",
+			test.rest.Start, test.rest.End, test.num), func(t *testing.T) {
+			r := test.rest
+
+			// Get the minimum size that a split restriction can be. Max size
+			// should be min + 1. This way we can check the size of each split.
+			num := test.num
+			if num <= 1 {
+				num = 1
+			}
+			min := (r.End - r.Start) / num
+
+			splits := r.EvenSplits(test.num)
+			prevEnd := r.Start
+			for _, split := range splits {
+				size := split.End - split.Start
+				// Check: Each restriction has at least 1 element.
+				if size == 0 {
+					t.Errorf("split restriction [%v, %v] is empty, size must be greater than 0.",
+						split.Start, split.End)
+				}
+				// Check: Restrictions are evenly split.
+				if size != min && size != min+1 {
+					t.Errorf("split restriction [%v, %v] has unexpected size. got: %v, want: %v or %v",
+						split.Start, split.End, size, min, min+1)
+				}
+				// Check: All elements are still in a split restrictions. This
+				// logic assumes that the splits are returned in order which
+				// isn't guaranteed by EvenSplits, but this check is way easier
+				// with the assumption.
+				if split.Start != prevEnd {
+					t.Errorf("restriction range [%v, %v] missing after splits.",
+						prevEnd, split.Start)
+				} else {
+					prevEnd = split.End
+				}
+			}
+			if prevEnd != r.End {
+				t.Errorf("restriction range [%v, %v] missing after splits.",
+					prevEnd, r.End)
+			}
+		})
+	}
+}
+
+// TestTracker_TryClaim validates both success and failure cases for TryClaim.
+func TestTracker_TryClaim(t *testing.T) {
+	// Test that TryClaim works as expected when called correctly.
+	t.Run("Correctness", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 1, 2, 3}},
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{15, 20, 50}},
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)
+					// If TryClaim returns false, check if there was an error.
+					if !ok && !rt.IsDone() {
+						t.Fatalf("tracker claiming %v failed, error: %v", pos, rt.GetError())
+					}
+				}
+			})
+		}
+	})
+
+	// Test that each invalid error case actually results in an error.
+	t.Run("Errors", func(t *testing.T) {
+		tests := []struct {
+			rest   Restriction
+			claims []int64
+		}{
+			// Claiming backwards.
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{0, 2, 1}},
+			// Claiming before start of restriction.
+			{rest: Restriction{Start: 10, End: 40}, claims: []int64{8}},
+			// Claiming after tracker signalled to stop.
+			{rest: Restriction{Start: 0, End: 3}, claims: []int64{4, 5}},
+		}
+		for _, test := range tests {
+			test := test
+			t.Run(fmt.Sprintf("(rest[%v, %v], claims = %v)",
+				test.rest.Start, test.rest.End, test.claims), func(t *testing.T) {
+				rt := NewTracker(test.rest)
+				for _, pos := range test.claims {
+					ok := rt.TryClaim(pos)
+					// Finish successfully if we got an error.
+					if !ok && !rt.IsDone() && rt.GetError() != nil {
+						return
+					}
+				}
+				t.Fatal("tracker did not fail on invalid claim")
+			})
+		}
+	})
+}
+
+// TestTracker_TrySplit tests that TrySplit follows its contract, meaning that
+// splits don't lose any elements, split fractions are clamped to 0 or 1, and
+// that TrySplit always splits at the nearest integer greater than the given
+// fraction.
+func TestTracker_TrySplit(t *testing.T) {
+	tests := []struct {
+		rest     Restriction
+		claimed  int64
+		fraction float64
+		splitPt  int64 // Index where we want the split to happen.

Review comment:
       Reading into the test further, it looks like "last element of primary" which I'm pretty sure is the correct thing, so this just means the comment should be updated to reflect that.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org