diff --git a/bar.go b/bar.go index e4278a2..19274d8 100644 --- a/bar.go +++ b/bar.go @@ -6,8 +6,10 @@ "fmt" "io" "io/ioutil" + "runtime" "strings" "sync" + "sync/atomic" "time" "unicode/utf8" @@ -46,6 +48,11 @@ done chan struct{} // shutdown is closed from master Progress goroutine only shutdown chan struct{} + + arbitraryCurrent struct { + lock uint32 + current int64 + } } type ( @@ -182,25 +189,6 @@ } } -// SetTotal sets total dynamically. -// Set final to true, when total is known, it will trigger bar complete event. -func (b *Bar) SetTotal(total int64, final bool) bool { - select { - case b.operateState <- func(s *bState) { - if total > 0 { - s.total = total - } - if final { - s.current = s.total - s.toComplete = true - } - }: - return true - case <-b.done: - return false - } -} - // SetRefill sets refill, if supported by underlying Filler. func (b *Bar) SetRefill(upto int) { b.operateState <- func(s *bState) { @@ -208,6 +196,39 @@ f.SetRefill(upto) } } +} + +// SetTotal sets total dynamically. +// Set final to true, when total is known, it will trigger bar complete event. +func (b *Bar) SetTotal(total int64, final bool) bool { + select { + case b.operateState <- func(s *bState) { + if total > 0 { + s.total = total + } + if final { + s.current = s.total + s.toComplete = true + } + }: + return true + case <-b.done: + return false + } +} + +// SetCurrent sets progress' current to arbitrary amount. +func (b *Bar) SetCurrent(current int64, wdd ...time.Duration) { + if current <= 0 { + return + } + for !atomic.CompareAndSwapUint32(&b.arbitraryCurrent.lock, 0, 1) { + runtime.Gosched() + } + last := b.arbitraryCurrent.current + b.IncrBy(int(current-last), wdd...) + b.arbitraryCurrent.current = current + atomic.StoreUint32(&b.arbitraryCurrent.lock, 0) } // Increment is a shorthand for b.IncrBy(1).