uploadprovider: allow closing used sources

Helps detecting when uploadprovider has finished
with the source. This avoids deadlock, should the
same reader be shared by multiple uploads that sync
with each other.

Signed-off-by: Tonis Tiigi <[email protected]>

Author: tonistiigi

frontend/dockerfile/dockerfile_test.go

@@ -5256,7 +5256,7 @@ COPY foo /
 	defer c.Close()
 
 	up := uploadprovider.New()
-	url := up.Add(buf)
+	url := up.Add(io.NopCloser(buf))
 
 	_, err = f.Solve(sb.Context(), c, client.SolveOpt{
 		FrontendAttrs: map[string]string{
@@ -5301,7 +5301,7 @@ COPY foo bar
 	defer c.Close()
 
 	up := uploadprovider.New()
-	url := up.Add(buf)
+	url := up.Add(io.NopCloser(buf))
 
 	// repeat with changed default args should match the old cache
 	destDir := t.TempDir()

session/upload/uploadprovider/provider.go

@@ -13,15 +13,15 @@ import (
 )
 
 func New() *Uploader {
-	return &Uploader{m: map[string]io.Reader{}}
+	return &Uploader{m: map[string]io.ReadCloser{}}
 }
 
 type Uploader struct {
 	mu sync.Mutex
-	m  map[string]io.Reader
+	m  map[string]io.ReadCloser
 }
 
-func (hp *Uploader) Add(r io.Reader) string {
+func (hp *Uploader) Add(r io.ReadCloser) string {
 	id := identity.NewID()
 	hp.m[id] = r
 	return "http://buildkit-session/" + id
@@ -51,6 +51,11 @@ func (hp *Uploader) Pull(stream upload.Upload_PullServer) error {
 	hp.mu.Unlock()
 
 	_, err := io.Copy(&writer{stream}, r)
+
+	if err1 := r.Close(); err == nil {
+		err = err1
+	}
+
 	return err
 }