gitpipe: Propagate context cancellation in object data pipelinepks-gitpipe-context-cancellation-errors

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 87f2d6b83..8d05e77f0 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