diff --git a/bar.go b/bar.go index 9472f5a..db8a350 100644 --- a/bar.go +++ b/bar.go @@ -90,7 +90,14 @@ if r == nil { panic("expected non nil io.Reader") } - return b.newProxyReader(r) + result := make(chan bool) + select { + case b.operateState <- func(s *bState) { result <- s.triggerComplete }: + triggerComplete := <-result + return b.newProxyReader(r, !triggerComplete) + case <-b.done: + return nil + } } // ID returs id of the bar. diff --git a/proxyreader.go b/proxyreader.go index 582a55e..c9a14b0 100644 --- a/proxyreader.go +++ b/proxyreader.go @@ -8,13 +8,14 @@ type proxyReader struct { io.ReadCloser - bar *Bar + bar *Bar + totalUnknown bool } func (x proxyReader) Read(p []byte) (int, error) { n, err := x.ReadCloser.Read(p) x.bar.IncrBy(n) - if err == io.EOF { + if x.totalUnknown && err == io.EOF { go x.bar.SetTotal(-1, true) } return n, err @@ -28,7 +29,7 @@ func (x proxyWriterTo) WriteTo(w io.Writer) (int64, error) { n, err := x.wt.WriteTo(w) x.bar.IncrInt64(n) - if err == io.EOF { + if x.totalUnknown && err == io.EOF { go x.bar.SetTotal(-1, true) } return n, err @@ -61,8 +62,8 @@ return n, err } -func (b *Bar) newProxyReader(r io.Reader) (rc io.ReadCloser) { - pr := proxyReader{toReadCloser(r), b} +func (b *Bar) newProxyReader(r io.Reader, totalUnknown bool) (rc io.ReadCloser) { + pr := proxyReader{toReadCloser(r), b, totalUnknown} if wt, ok := r.(io.WriterTo); ok { pw := proxyWriterTo{pr, wt} if b.hasEwma {