diff options
-rw-r--r-- | internal/git/gitpipe/catfile_object.go | 3 | ||||
-rw-r--r-- | internal/git/gitpipe/catfile_object_iterator.go | 38 | ||||
-rw-r--r-- | internal/git/gitpipe/catfile_object_test.go | 7 | ||||
-rw-r--r-- | internal/git/gitpipe/pipeline_test.go | 2 |
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 |