[refactor] Use plain loops with iterators (jaegertracing#6722)

## Description of the changes - Follow-up to jaegertracing#6714 ## How was this change tested? - CI Signed-off-by: Yuri Shkuro <[email protected]>
zzzk1 · Feb 13, 2025 · 5182729 · 5182729
1 parent 0bc62de
commit 5182729
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 60 deletions.
diff --git a/cmd/query/app/apiv3/grpc_handler.go b/cmd/query/app/apiv3/grpc_handler.go
@@ -131,21 +131,17 @@ func receiveTraces(
 	seq iter.Seq2[[]ptrace.Traces, error],
 	sendFn func(*api_v3.TracesData) error,
 ) error {
-	var capturedErr error
-	seq(func(traces []ptrace.Traces, err error) bool {
+	for traces, err := range seq {
 		if err != nil {
-			capturedErr = err
-			return false
+			return err
 		}
 		for _, trace := range traces {
 			tracesData := api_v3.TracesData(trace)
 			if err := sendFn(&tracesData); err != nil {
-				capturedErr = status.Error(codes.Internal,
+				return status.Error(codes.Internal,
 					fmt.Sprintf("failed to send response stream chunk to client: %v", err))
-				return false
 			}
 		}
-		return true
-	})
-	return capturedErr
+	}
+	return nil
 }
diff --git a/internal/jiter/iter.go b/internal/jiter/iter.go
@@ -10,34 +10,22 @@ import (
 
 func CollectWithErrors[V any](seq iter.Seq2[V, error]) ([]V, error) {
 	var result []V
-	var err error
-	seq(func(v V, e error) bool {
-		if e != nil {
-			err = e
-			return false
+	for v, err := range seq {
+		if err != nil {
+			return nil, err
 		}
 		result = append(result, v)
-		return true
-	})
-	if err != nil {
-		return nil, err
 	}
 	return result, nil
 }
 
 func FlattenWithErrors[V any](seq iter.Seq2[[]V, error]) ([]V, error) {
 	var result []V
-	var err error
-	seq(func(v []V, e error) bool {
-		if e != nil {
-			err = e
-			return false
+	for v, err := range seq {
+		if err != nil {
+			return nil, err
 		}
 		result = append(result, v...)
-		return true
-	})
-	if err != nil {
-		return nil, err
 	}
 	return result, nil
 }
diff --git a/internal/jiter/iter_test.go b/internal/jiter/iter_test.go
@@ -4,60 +4,71 @@
 package jiter
 
 import (
-	"iter"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestCollectWithErrors(t *testing.T) {
+	type item struct {
+		str string
+		err error
+	}
 	tests := []struct {
 		name     string
-		seq      iter.Seq2[string, error]
+		items    []item
 		expected []string
 		err      error
 	}{
 		{
 			name: "no errors",
-			seq: func(yield func(string, error) bool) {
-				yield("a", nil)
-				yield("b", nil)
-				yield("c", nil)
+			items: []item{
+				{str: "a", err: nil},
+				{str: "b", err: nil},
+				{str: "c", err: nil},
 			},
 			expected: []string{"a", "b", "c"},
 		},
 		{
 			name: "first error",
-			seq: func(yield func(string, error) bool) {
-				yield("a", nil)
-				yield("b", nil)
-				yield("c", assert.AnError)
+			items: []item{
+				{str: "a", err: nil},
+				{str: "b", err: nil},
+				{str: "c", err: assert.AnError},
 			},
 			err: assert.AnError,
 		},
 		{
 			name: "second error",
-			seq: func(yield func(string, error) bool) {
-				yield("a", nil)
-				yield("b", assert.AnError)
-				yield("c", nil)
+			items: []item{
+				{str: "a", err: nil},
+				{str: "b", err: assert.AnError},
+				{str: "c", err: nil},
 			},
 			err: assert.AnError,
 		},
 		{
 			name: "third error",
-			seq: func(yield func(string, error) bool) {
-				yield("a", nil)
-				yield("b", nil)
-				yield("c", assert.AnError)
+			items: []item{
+				{str: "a", err: nil},
+				{str: "b", err: nil},
+				{str: "c", err: assert.AnError},
 			},
+
 			err: assert.AnError,
 		},
 	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			result, err := CollectWithErrors(test.seq)
+			seq := func(yield func(string, error) bool) {
+				for _, item := range test.items {
+					if !yield(item.str, item.err) {
+						return
+					}
+				}
+			}
+			result, err := CollectWithErrors(seq)
 			if test.err != nil {
 				require.ErrorIs(t, err, test.err)
 			} else {
@@ -69,43 +80,54 @@ func TestCollectWithErrors(t *testing.T) {
 }
 
 func TestFlattenWithErrors(t *testing.T) {
+	type item struct {
+		strs []string
+		err  error
+	}
 	tests := []struct {
 		name     string
-		seq      iter.Seq2[[]string, error]
+		items    []item
 		expected []string
 		err      error
 	}{
 		{
 			name: "no errors",
-			seq: func(yield func([]string, error) bool) {
-				yield([]string{"a", "b", "c"}, nil)
-				yield([]string{"d", "e", "f"}, nil)
-				yield([]string{"g", "h", "i"}, nil)
+			items: []item{
+				{strs: []string{"a", "b", "c"}, err: nil},
+				{strs: []string{"d", "e", "f"}, err: nil},
+				{strs: []string{"g", "h", "i"}, err: nil},
 			},
 			expected: []string{"a", "b", "c", "d", "e", "f", "g", "h", "i"},
 		},
 		{
 			name: "first error",
-			seq: func(yield func([]string, error) bool) {
-				yield([]string{"a", "b", "c"}, nil)
-				yield([]string{"d", "e", "f"}, assert.AnError)
-				yield([]string{"g", "h", "i"}, nil)
+			items: []item{
+				{strs: []string{"a", "b", "c"}, err: nil},
+				{strs: []string{"d", "e", "f"}, err: assert.AnError},
+				{strs: []string{"g", "h", "i"}, err: nil},
 			},
 			err: assert.AnError,
 		},
 		{
 			name: "second error",
-			seq: func(yield func([]string, error) bool) {
-				yield([]string{"a", "b", "c"}, nil)
-				yield([]string{"d", "e", "f"}, nil)
-				yield([]string{"g", "h", "i"}, assert.AnError)
+			items: []item{
+				{strs: []string{"a", "b", "c"}, err: nil},
+				{strs: []string{"d", "e", "f"}, err: nil},
+				{strs: []string{"g", "h", "i"}, err: assert.AnError},
 			},
 			err: assert.AnError,
 		},
 	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			result, err := FlattenWithErrors(test.seq)
+			seq := func(yield func([]string, error) bool) {
+				for _, item := range test.items {
+					if !yield(item.strs, item.err) {
+						return
+					}
+				}
+			}
+			result, err := FlattenWithErrors(seq)
 			if test.err != nil {
 				require.ErrorIs(t, err, test.err)
 			} else {