Browse code
Fix env.WriteTo count return
Some calls like json.Encoder.Encode mask the number of bytes written to
an io.Writer. The solution provides a wrapper io.Writer around the
actual io.Writer that allows multiple calls to Write to be considered as
one and allow access to this count.
Signed-off-by: Pierre Wacrenier <pierre.wacrenier@gmail.com>
Pierre Wacrenier authored on 2015/01/19 09:27:14Showing 3 changed files
| ... | ... |
@@ -7,6 +7,8 @@ import ( |
| 7 | 7 |
"io" |
| 8 | 8 |
"strconv" |
| 9 | 9 |
"strings" |
| 10 |
+ |
|
| 11 |
+ "github.com/docker/docker/utils" |
|
| 10 | 12 |
) |
| 11 | 13 |
|
| 12 | 14 |
type Env []string |
| ... | ... |
@@ -242,9 +244,10 @@ func (env *Env) Encode(dst io.Writer) error {
|
| 242 | 242 |
return nil |
| 243 | 243 |
} |
| 244 | 244 |
|
| 245 |
-func (env *Env) WriteTo(dst io.Writer) (n int64, err error) {
|
|
| 246 |
- // FIXME: return the number of bytes written to respect io.WriterTo |
|
| 247 |
- return 0, env.Encode(dst) |
|
| 245 |
+func (env *Env) WriteTo(dst io.Writer) (int64, error) {
|
|
| 246 |
+ wc := utils.NewWriteCounter(dst) |
|
| 247 |
+ err := env.Encode(wc) |
|
| 248 |
+ return wc.Count, err |
|
| 248 | 249 |
} |
| 249 | 250 |
|
| 250 | 251 |
func (env *Env) Import(src interface{}) (err error) {
|
| ... | ... |
@@ -516,3 +516,24 @@ func ReadDockerIgnore(path string) ([]string, error) {
|
| 516 | 516 |
} |
| 517 | 517 |
return excludes, nil |
| 518 | 518 |
} |
| 519 |
+ |
|
| 520 |
+// Wrap a concrete io.Writer and hold a count of the number |
|
| 521 |
+// of bytes written to the writer during a "session". |
|
| 522 |
+// This can be convenient when write return is masked |
|
| 523 |
+// (e.g., json.Encoder.Encode()) |
|
| 524 |
+type WriteCounter struct {
|
|
| 525 |
+ Count int64 |
|
| 526 |
+ Writer io.Writer |
|
| 527 |
+} |
|
| 528 |
+ |
|
| 529 |
+func NewWriteCounter(w io.Writer) *WriteCounter {
|
|
| 530 |
+ return &WriteCounter{
|
|
| 531 |
+ Writer: w, |
|
| 532 |
+ } |
|
| 533 |
+} |
|
| 534 |
+ |
|
| 535 |
+func (wc *WriteCounter) Write(p []byte) (count int, err error) {
|
|
| 536 |
+ count, err = wc.Writer.Write(p) |
|
| 537 |
+ wc.Count += int64(count) |
|
| 538 |
+ return |
|
| 539 |
+} |
| ... | ... |
@@ -1,7 +1,9 @@ |
| 1 | 1 |
package utils |
| 2 | 2 |
|
| 3 | 3 |
import ( |
| 4 |
+ "bytes" |
|
| 4 | 5 |
"os" |
| 6 |
+ "strings" |
|
| 5 | 7 |
"testing" |
| 6 | 8 |
) |
| 7 | 9 |
|
| ... | ... |
@@ -97,3 +99,26 @@ func TestReadSymlinkedDirectoryToFile(t *testing.T) {
|
| 97 | 97 |
t.Errorf("failed to remove symlink: %s", err)
|
| 98 | 98 |
} |
| 99 | 99 |
} |
| 100 |
+ |
|
| 101 |
+func TestWriteCounter(t *testing.T) {
|
|
| 102 |
+ dummy1 := "This is a dummy string." |
|
| 103 |
+ dummy2 := "This is another dummy string." |
|
| 104 |
+ totalLength := int64(len(dummy1) + len(dummy2)) |
|
| 105 |
+ |
|
| 106 |
+ reader1 := strings.NewReader(dummy1) |
|
| 107 |
+ reader2 := strings.NewReader(dummy2) |
|
| 108 |
+ |
|
| 109 |
+ var buffer bytes.Buffer |
|
| 110 |
+ wc := NewWriteCounter(&buffer) |
|
| 111 |
+ |
|
| 112 |
+ reader1.WriteTo(wc) |
|
| 113 |
+ reader2.WriteTo(wc) |
|
| 114 |
+ |
|
| 115 |
+ if wc.Count != totalLength {
|
|
| 116 |
+ t.Errorf("Wrong count: %d vs. %d", wc.Count, totalLength)
|
|
| 117 |
+ } |
|
| 118 |
+ |
|
| 119 |
+ if buffer.String() != dummy1+dummy2 {
|
|
| 120 |
+ t.Error("Wrong message written")
|
|
| 121 |
+ } |
|
| 122 |
+} |