gitpipe: Propagate context cancellation in object data pipeline

When the context gets cancelled while we're iterating over results from
the object data 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, 38 insertions, 12 deletions

diff --git a/internal/git/gitpipe/catfile_object.go b/internal/git/gitpipe/catfile_object.go
index a393ed2ae..f23496158 100644
--- a/internal/git/gitpipe/catfile_object.go
+++ b/internal/git/gitpipe/catfile_object.go
@@ -168,7 +168,8 @@ func CatfileObject(
 	}()
 
 	return &catfileObjectIterator{
-		ch: resultChan,
+		ctx: ctx,
+		ch:  resultChan,
 	}, nil
 }
 
diff --git a/internal/git/gitpipe/catfile_object_iterator.go b/internal/git/gitpipe/catfile_object_iterator.go
index c457088e0..05b00e3e8 100644
--- a/internal/git/gitpipe/catfile_object_iterator.go
+++ b/internal/git/gitpipe/catfile_object_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"
+)
 
 // CatfileObjectIterator is an iterator returned by the Revlist function.
 type CatfileObjectIterator interface {
@@ -10,7 +14,7 @@ type CatfileObjectIterator interface {
 }
 
 // NewCatfileObjectIterator returns a new CatfileObjectIterator for the given items.
-func NewCatfileObjectIterator(items []CatfileObjectResult) CatfileObjectIterator {
+func NewCatfileObjectIterator(ctx context.Context, items []CatfileObjectResult) CatfileObjectIterator {
 	itemChan := make(chan CatfileObjectResult, len(items))
 	for _, item := range items {
 		itemChan <- item
@@ -18,11 +22,13 @@ func NewCatfileObjectIterator(items []CatfileObjectResult) CatfileObjectIterator
 	close(itemChan)
 
 	return &catfileObjectIterator{
-		ch: itemChan,
+		ctx: ctx,
+		ch:  itemChan,
 	}
 }
 
 type catfileObjectIterator struct {
+	ctx    context.Context
 	ch     <-chan CatfileObjectResult
 	result CatfileObjectResult
 }
@@ -32,13 +38,31 @@ func (it *catfileObjectIterator) 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 = CatfileObjectResult{err: it.ctx.Err()}
 		return false
+	default:
 	}
 
-	return true
+	select {
+	case <-it.ctx.Done():
+		it.result = CatfileObjectResult{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 *catfileObjectIterator) Err() error {
diff --git a/internal/git/gitpipe/catfile_object_test.go b/internal/git/gitpipe/catfile_object_test.go
index 17f07f1f8..bf96a19e0 100644
--- a/internal/git/gitpipe/catfile_object_test.go
+++ b/internal/git/gitpipe/catfile_object_test.go
@@ -144,8 +144,9 @@ func TestCatfileObject(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, CatfileObjectResult{}, it.Result())
+		require.Equal(t, context.Canceled, it.Err())
+		require.Equal(t, CatfileObjectResult{
+			err: context.Canceled,
+		}, it.Result())
 	})
 }
diff --git a/internal/git/gitpipe/pipeline_test.go b/internal/git/gitpipe/pipeline_test.go
index 544ad552d..3d34db488 100644
--- a/internal/git/gitpipe/pipeline_test.go
+++ b/internal/git/gitpipe/pipeline_test.go
@@ -310,7 +310,7 @@ func TestPipeline_revlist(t *testing.T) {
 			}
 		}
 
-		require.NoError(t, catfileObjectIter.Err())
+		require.Equal(t, context.Canceled, catfileObjectIter.Err())
 
 		// Context cancellation is timing sensitive: at the point of cancelling the context,
 		// the last pipeline step may already have queued up an additional result. We thus