diff --git a/bar.go b/bar.go index 043163d..1a4c66f 100644 --- a/bar.go +++ b/bar.go @@ -5,7 +5,6 @@ "context" "fmt" "io" - "io/ioutil" "log" "strings" "time" @@ -34,16 +33,15 @@ priority int // used by heap index int // used by heap - extendedLines int - toShutdown bool - toDrop bool - noPop bool - hasAverageDecorators bool - hasEwmaDecorators bool - operateState chan func(*bState) - frameCh chan io.Reader - syncTableCh chan [][]chan int - completed chan bool + extendedLines int + toShutdown bool + toDrop bool + noPop bool + hasEwmaDecorators bool + operateState chan func(*bState) + frameCh chan io.Reader + syncTableCh chan [][]chan int + completed chan bool // cancel is called either by user or on complete event cancel func() @@ -113,19 +111,12 @@ } // ProxyReader wraps r with metrics required for progress tracking. +// Panics if r is nil. func (b *Bar) ProxyReader(r io.Reader) io.ReadCloser { if r == nil { - return nil - } - rc, ok := r.(io.ReadCloser) - if !ok { - rc = ioutil.NopCloser(r) - } - prox := &proxyReader{rc, b, time.Now()} - if wt, ok := r.(io.WriterTo); ok { - return &proxyWriterTo{prox, wt} - } - return prox + panic("expected non nil io.Reader") + } + return newProxyReader(r, b) } // ID returs id of the bar. @@ -248,9 +239,6 @@ // iteration's duration. Panics if called before *Bar.Incr... family // methods. func (b *Bar) DecoratorEwmaUpdate(dur time.Duration) { - if !b.hasEwmaDecorators { - return - } select { case b.operateState <- func(s *bState) { ewmaIterationUpdate(false, s, dur) @@ -264,9 +252,6 @@ // if you need to adjust start time of all average based decorators // or after progress resume. func (b *Bar) DecoratorAverageAdjust(start time.Time) { - if !b.hasAverageDecorators { - return - } select { case b.operateState <- func(s *bState) { for _, d := range s.averageDecorators { @@ -389,7 +374,6 @@ s.ewmaDecorators = ewmaDecorators s.shutdownListeners = shutdownListeners } - b.hasAverageDecorators = len(averageDecorators) != 0 b.hasEwmaDecorators = len(ewmaDecorators) != 0 } diff --git a/proxyreader.go b/proxyreader.go index 8323986..0efb09a 100644 --- a/proxyreader.go +++ b/proxyreader.go @@ -2,42 +2,90 @@ import ( "io" + "io/ioutil" "time" ) type proxyReader struct { io.ReadCloser bar *Bar - iT time.Time } -func (prox *proxyReader) Read(p []byte) (int, error) { - n, err := prox.ReadCloser.Read(p) - if n > 0 { - prox.bar.IncrBy(n) - prox.bar.DecoratorEwmaUpdate(time.Since(prox.iT)) - prox.iT = time.Now() - } +func (x *proxyReader) Read(p []byte) (int, error) { + n, err := x.ReadCloser.Read(p) + x.bar.IncrBy(n) if err == io.EOF { - go prox.bar.SetTotal(0, true) + go x.bar.SetTotal(0, true) } return n, err } type proxyWriterTo struct { - *proxyReader - wt io.WriterTo + io.ReadCloser // *proxyReader + wt io.WriterTo + bar *Bar } -func (prox *proxyWriterTo) WriteTo(w io.Writer) (int64, error) { - n, err := prox.wt.WriteTo(w) - if n > 0 { - prox.bar.IncrInt64(n) - prox.bar.DecoratorEwmaUpdate(time.Since(prox.iT)) - prox.iT = time.Now() - } +func (x *proxyWriterTo) WriteTo(w io.Writer) (int64, error) { + n, err := x.wt.WriteTo(w) + x.bar.IncrInt64(n) if err == io.EOF { - go prox.bar.SetTotal(0, true) + go x.bar.SetTotal(0, true) } return n, err } + +type ewmaProxyReader struct { + io.ReadCloser // *proxyReader + bar *Bar + iT time.Time +} + +func (x *ewmaProxyReader) Read(p []byte) (int, error) { + n, err := x.ReadCloser.Read(p) + if n > 0 { + x.bar.DecoratorEwmaUpdate(time.Since(x.iT)) + x.iT = time.Now() + } + return n, err +} + +type ewmaProxyWriterTo struct { + io.ReadCloser // *ewmaProxyReader + wt io.WriterTo // *proxyWriterTo + bar *Bar + iT time.Time +} + +func (x *ewmaProxyWriterTo) WriteTo(w io.Writer) (int64, error) { + n, err := x.wt.WriteTo(w) + if n > 0 { + x.bar.DecoratorEwmaUpdate(time.Since(x.iT)) + x.iT = time.Now() + } + return n, err +} + +func newProxyReader(r io.Reader, bar *Bar) io.ReadCloser { + wt, isWriterTo := r.(io.WriterTo) + rc := toReadCloser(r) + rc = &proxyReader{rc, bar} + + if bar.hasEwmaDecorators { + now := time.Now() + rc = &ewmaProxyReader{rc, bar, now} + if isWriterTo { + rc = &ewmaProxyWriterTo{rc, wt, bar, now} + } + } else if isWriterTo { + rc = &proxyWriterTo{rc, wt, bar} + } + return rc +} + +func toReadCloser(r io.Reader) io.ReadCloser { + if rc, ok := r.(io.ReadCloser); ok { + return rc + } + return ioutil.NopCloser(r) +}