diff --git a/internal/handler/stream_miss_test.go b/internal/handler/stream_miss_test.go index 172be1bc..8ffc3699 100644 --- a/internal/handler/stream_miss_test.go +++ b/internal/handler/stream_miss_test.go @@ -530,12 +530,20 @@ func TestCacheMissSlowClientFailsWithoutStoredCopy(t *testing.T) { lowerFeedLimit(t, 64<<10) proxy, _, store, _ := setupTestProxy(t) store.storeErr = errors.New("no space left on device") - proxy.Fetcher = &countingFetcher{content: streamMissContent} + fetcher := &closeSignalFetcher{countingFetcher: countingFetcher{content: streamMissContent}, closed: make(chan struct{})} + proxy.Fetcher = fetcher res, err := proxy.GetOrFetchArtifactFromURL(StreamCacheMiss(context.Background()), "npm", "big", "1.0.0", "big-1.0.0.tgz", "https://example.com/big.tgz") if err != nil { t.Fatal(err) } + // Read nothing until the upstream copy is over, so the client is behind + // rather than keeping up, which it can on a slow runner. + select { + case <-fetcher.closed: + case <-time.After(10 * time.Second): + t.Fatal("upstream body never closed") + } body, err := io.ReadAll(res.Reader) _ = res.Reader.Close() if err == nil { @@ -543,6 +551,36 @@ func TestCacheMissSlowClientFailsWithoutStoredCopy(t *testing.T) { } } +// closeSignalFetcher closes closed when the proxy closes the upstream body. +type closeSignalFetcher struct { + countingFetcher + closed chan struct{} + once sync.Once +} + +func (f *closeSignalFetcher) Fetch(ctx context.Context, url string) (*fetch.Artifact, error) { + return f.FetchWithHeaders(ctx, url, nil) +} + +func (f *closeSignalFetcher) FetchWithHeaders(ctx context.Context, url string, h http.Header) (*fetch.Artifact, error) { + a, err := f.countingFetcher.FetchWithHeaders(ctx, url, h) + if err != nil { + return nil, err + } + a.Body = &signalCloser{Reader: a.Body, close: func() { f.once.Do(func() { close(f.closed) }) }} + return a, nil +} + +type signalCloser struct { + io.Reader + close func() +} + +func (s *signalCloser) Close() error { + s.close() + return nil +} + func TestMissFeedCloseWhileWaiting(t *testing.T) { m := newMissFeed(4) done := make(chan error, 1)