Merge pull request #13512 from moysesb/progressreader_close_silently

Close progressreader silently when read is complete.

pkg/progressreader/progressreader.go

@@ -23,6 +23,7 @@ type Config struct {
 func New(newReader Config) *Config {
 	return &newReader
 }
+
 func (config *Config) Read(p []byte) (n int, err error) {
 	read, err := config.In.Read(p)
 	config.Current += read
@@ -34,17 +35,30 @@ func (config *Config) Read(p []byte) (n int, err error) {
 		}
 	}
 	if config.Current-config.LastUpdate > updateEvery || err != nil {
-		config.Out.Write(config.Formatter.FormatProgress(config.ID, config.Action, &jsonmessage.JSONProgress{Current: config.Current, Total: config.Size}))
+		updateProgress(config)
 		config.LastUpdate = config.Current
 	}
-	// Send newline when complete
-	if config.NewLines && err != nil && read == 0 {
-		config.Out.Write(config.Formatter.FormatStatus("", ""))
+
+	if err != nil && read == 0 {
+		updateProgress(config)
+		if config.NewLines {
+			config.Out.Write(config.Formatter.FormatStatus("", ""))
+		}
 	}
 	return read, err
 }
+
 func (config *Config) Close() error {
-	config.Current = config.Size
-	config.Out.Write(config.Formatter.FormatProgress(config.ID, config.Action, &jsonmessage.JSONProgress{Current: config.Current, Total: config.Size}))
+	if config.Current < config.Size {
+		//print a full progress bar when closing prematurely
+		config.Current = config.Size
+		updateProgress(config)
+	}
 	return config.In.Close()
 }
+
+func updateProgress(config *Config) {
+	progress := jsonmessage.JSONProgress{Current: config.Current, Total: config.Size}
+	fmtMessage := config.Formatter.FormatProgress(config.ID, config.Action, &progress)
+	config.Out.Write(fmtMessage)
+}

pkg/progressreader/progressreader_test.go

@@ -0,0 +1,94 @@
+package progressreader
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"io/ioutil"
+	"testing"
+
+	"github.com/docker/docker/pkg/streamformatter"
+)
+
+func TestOutputOnPrematureClose(t *testing.T) {
+	var outBuf bytes.Buffer
+	content := []byte("TESTING")
+	reader := ioutil.NopCloser(bytes.NewReader(content))
+	writer := bufio.NewWriter(&outBuf)
+
+	prCfg := Config{
+		In:        reader,
+		Out:       writer,
+		Formatter: streamformatter.NewStreamFormatter(),
+		Size:      len(content),
+		NewLines:  true,
+		ID:        "Test",
+		Action:    "Read",
+	}
+	pr := New(prCfg)
+
+	part := make([]byte, 4, 4)
+	_, err := io.ReadFull(pr, part)
+	if err != nil {
+		pr.Close()
+		t.Fatal(err)
+	}
+
+	if err := writer.Flush(); err != nil {
+		pr.Close()
+		t.Fatal(err)
+	}
+
+	tlen := outBuf.Len()
+	pr.Close()
+	if err := writer.Flush(); err != nil {
+		t.Fatal(err)
+	}
+
+	if outBuf.Len() == tlen {
+		t.Fatalf("Expected some output when closing prematurely")
+	}
+}
+
+func TestCompleteSilently(t *testing.T) {
+	var outBuf bytes.Buffer
+	content := []byte("TESTING")
+	reader := ioutil.NopCloser(bytes.NewReader(content))
+	writer := bufio.NewWriter(&outBuf)
+
+	prCfg := Config{
+		In:        reader,
+		Out:       writer,
+		Formatter: streamformatter.NewStreamFormatter(),
+		Size:      len(content),
+		NewLines:  true,
+		ID:        "Test",
+		Action:    "Read",
+	}
+	pr := New(prCfg)
+
+	out, err := ioutil.ReadAll(pr)
+	if err != nil {
+		pr.Close()
+		t.Fatal(err)
+	}
+	if string(out) != "TESTING" {
+		pr.Close()
+		t.Fatalf("Unexpected output %q from reader", string(out))
+	}
+
+	if err := writer.Flush(); err != nil {
+		pr.Close()
+		t.Fatal(err)
+	}
+
+	tlen := outBuf.Len()
+	pr.Close()
+	if err := writer.Flush(); err != nil {
+		t.Fatal(err)
+	}
+
+	if outBuf.Len() > tlen {
+		t.Fatalf("Should have closed silently when read is complete")
+	}
+}