gitpipe: Propagate context cancellation in object info pipeline

When the context gets cancelled while we're iterating over results from
the object info pipeline, then the iterator doesn't return the context
cancellation error to the caller when calling `iter.Err()`. It is thus
easy to assume at the calling side that the iterator has just finished
successfully and that there are no more results, while in reality we
only got a partial set of results.

Fix this issue by propagating context cancellation errors to the caller.
This fixes RPCs based on this pipeline to not return `OK` when there
indeed was an error.

Changelog: fixed

4 files changed, 45 insertions, 17 deletions

diff --git a/internal/git/gitpipe/catfile_info.go b/internal/git/gitpipe/catfile_info.go
index 314538fdf..47e2eabdd 100644
--- a/internal/git/gitpipe/catfile_info.go
+++ b/internal/git/gitpipe/catfile_info.go
@@ -140,7 +140,8 @@ func CatfileInfo(
 	}()
 
 	return &catfileInfoIterator{
-		ch: resultChan,
+		ctx: ctx,
+		ch:  resultChan,
 	}, nil
 }
 
@@ -215,7 +216,8 @@ func CatfileInfoAllObjects(
 	}()
 
 	return &catfileInfoIterator{
-		ch: resultChan,
+		ctx: ctx,
+		ch:  resultChan,
 	}
 }
 
diff --git a/internal/git/gitpipe/catfile_info_iterator.go b/internal/git/gitpipe/catfile_info_iterator.go
index 35c782699..c4403af5c 100644
--- a/internal/git/gitpipe/catfile_info_iterator.go
+++ b/internal/git/gitpipe/catfile_info_iterator.go
@@ -1,6 +1,10 @@
 package gitpipe
 
-import "gitlab.com/gitlab-org/gitaly/v14/internal/git"
+import (
+	"context"
+
+	"gitlab.com/gitlab-org/gitaly/v14/internal/git"
+)
 
 // CatfileInfoIterator is an iterator returned by the Revlist function.
 type CatfileInfoIterator interface {
@@ -10,7 +14,7 @@ type CatfileInfoIterator interface {
 }
 
 // NewCatfileInfoIterator returns a new CatfileInfoIterator for the given items.
-func NewCatfileInfoIterator(items []CatfileInfoResult) CatfileInfoIterator {
+func NewCatfileInfoIterator(ctx context.Context, items []CatfileInfoResult) CatfileInfoIterator {
 	itemChan := make(chan CatfileInfoResult, len(items))
 	for _, item := range items {
 		itemChan <- item
@@ -18,11 +22,13 @@ func NewCatfileInfoIterator(items []CatfileInfoResult) CatfileInfoIterator {
 	close(itemChan)
 
 	return &catfileInfoIterator{
-		ch: itemChan,
+		ctx: ctx,
+		ch:  itemChan,
 	}
 }
 
 type catfileInfoIterator struct {
+	ctx    context.Context
 	ch     <-chan CatfileInfoResult
 	result CatfileInfoResult
 }
@@ -32,13 +38,31 @@ func (it *catfileInfoIterator) Next() bool {
 		return false
 	}
 
-	var ok bool
-	it.result, ok = <-it.ch
-	if !ok || it.result.err != nil {
+	// Prioritize context cancellation errors so that we don't try to fetch results anymore when
+	// the context is done.
+	select {
+	case <-it.ctx.Done():
+		it.result = CatfileInfoResult{err: it.ctx.Err()}
 		return false
+	default:
 	}
 
-	return true
+	select {
+	case <-it.ctx.Done():
+		it.result = CatfileInfoResult{err: it.ctx.Err()}
+		return false
+	case result, ok := <-it.ch:
+		if !ok {
+			return false
+		}
+
+		it.result = result
+		if result.err != nil {
+			return false
+		}
+
+		return true
+	}
 }
 
 func (it *catfileInfoIterator) Err() error {
diff --git a/internal/git/gitpipe/catfile_info_test.go b/internal/git/gitpipe/catfile_info_test.go
index 8df257213..d2f7c2b37 100644
--- a/internal/git/gitpipe/catfile_info_test.go
+++ b/internal/git/gitpipe/catfile_info_test.go
@@ -182,9 +182,10 @@ func TestCatfileInfo(t *testing.T) {
 		cancel()
 
 		require.False(t, it.Next())
-		// This is a bug: we expect to get the cancelled context here.
-		require.NoError(t, it.Err())
-		require.Equal(t, CatfileInfoResult{}, it.Result())
+		require.Equal(t, context.Canceled, it.Err())
+		require.Equal(t, CatfileInfoResult{
+			err: context.Canceled,
+		}, it.Result())
 	})
 }
 
@@ -233,8 +234,9 @@ func TestCatfileInfoAllObjects(t *testing.T) {
 		cancel()
 
 		require.False(t, it.Next())
-		// This is a bug: we expect to get the cancelled context here.
-		require.NoError(t, it.Err())
-		require.Equal(t, CatfileInfoResult{}, it.Result())
+		require.Equal(t, context.Canceled, it.Err())
+		require.Equal(t, CatfileInfoResult{
+			err: context.Canceled,
+		}, it.Result())
 	})
 }
diff --git a/internal/git/gitpipe/catfile_object_test.go b/internal/git/gitpipe/catfile_object_test.go
index b847a6a75..17f07f1f8 100644
--- a/internal/git/gitpipe/catfile_object_test.go
+++ b/internal/git/gitpipe/catfile_object_test.go
@@ -80,7 +80,7 @@ func TestCatfileObject(t *testing.T) {
 			require.NoError(t, err)
 			defer cancel()
 
-			it, err := CatfileObject(ctx, objectReader, NewCatfileInfoIterator(tc.catfileInfoInputs))
+			it, err := CatfileObject(ctx, objectReader, NewCatfileInfoIterator(ctx, tc.catfileInfoInputs))
 			require.NoError(t, err)
 
 			var results []CatfileObjectResult
@@ -128,7 +128,7 @@ func TestCatfileObject(t *testing.T) {
 		require.NoError(t, err)
 		defer objectReaderCancel()
 
-		it, err := CatfileObject(ctx, objectReader, NewCatfileInfoIterator([]CatfileInfoResult{
+		it, err := CatfileObject(ctx, objectReader, NewCatfileInfoIterator(ctx, []CatfileInfoResult{
 			{ObjectInfo: &catfile.ObjectInfo{Oid: lfsPointer1, Type: "blob", Size: 133}},
 			{ObjectInfo: &catfile.ObjectInfo{Oid: lfsPointer2, Type: "blob", Size: 127}},
 		}))