diff --git a/bar.go b/bar.go index a9131a0..7d83268 100644 --- a/bar.go +++ b/bar.go @@ -338,7 +338,7 @@ // priority, i.e. bar will be on top. If you don't need to set priority // dynamically, better use BarPriority option. func (b *Bar) SetPriority(priority int) { - b.container.UpdateBarPriority(b, priority) + b.container.UpdateBarPriority(b, priority, false) } // Abort interrupts bar's running goroutine. Abort won't be engaged diff --git a/heap_manager.go b/heap_manager.go index 2f07664..b8d2054 100644 --- a/heap_manager.go +++ b/heap_manager.go @@ -36,6 +36,7 @@ type fixData struct { bar *Bar priority int + lazy bool } func (m heapManager) run() { @@ -96,12 +97,13 @@ close(data.iter) case h_fix: data := req.data.(fixData) - bar, priority := data.bar, data.priority - if bar.index < 0 { + if data.bar.index < 0 { break } - bar.priority = priority - heap.Fix(&bHeap, bar.index) + data.bar.priority = data.priority + if !data.lazy { + heap.Fix(&bHeap, data.bar.index) + } case h_state: ch := req.data.(chan<- bool) ch <- sync || l != bHeap.Len() @@ -136,8 +138,8 @@ m <- heapRequest{cmd: h_drain, data: data} } -func (m heapManager) fix(b *Bar, priority int) { - data := fixData{b, priority} +func (m heapManager) fix(b *Bar, priority int, lazy bool) { + data := fixData{b, priority, lazy} m <- heapRequest{cmd: h_fix, data: data} } diff --git a/progress.go b/progress.go index cc4e3e1..f275be3 100644 --- a/progress.go +++ b/progress.go @@ -201,10 +201,15 @@ } } -// UpdateBarPriority same as *Bar.SetPriority(int). -func (p *Progress) UpdateBarPriority(b *Bar, priority int) { +// UpdateBarPriority either immediately or lazy. +// With lazy flag order is updated after the next refresh cycle. +// If you don't care about laziness just use *Bar.SetPriority(int). +func (p *Progress) UpdateBarPriority(b *Bar, priority int, lazy bool) { + if b == nil { + return + } select { - case p.operateState <- func(s *pState) { s.hm.fix(b, priority) }: + case p.operateState <- func(s *pState) { s.hm.fix(b, priority, lazy) }: case <-p.done: } } diff --git a/progress_test.go b/progress_test.go index 01156b5..2de83af 100644 --- a/progress_test.go +++ b/progress_test.go @@ -208,7 +208,7 @@ shutdown := make(chan interface{}) ctx, cancel := context.WithCancel(context.Background()) p := mpb.NewWithContext(ctx, - mpb.WithOutput(io.Discard), + mpb.WithOutput(io.Discard), // auto refresh is disabled because of io.Discard mpb.WithShutdownNotifier(shutdown), ) a := p.AddBar(100, mpb.BarPriority(1)) @@ -221,10 +221,58 @@ c: "c", } - p.UpdateBarPriority(c, 2) - p.UpdateBarPriority(b, 3) - - cancel() + p.UpdateBarPriority(c, 2, false) + p.UpdateBarPriority(b, 3, false) + + go cancel() + + bars := (<-shutdown).([]*mpb.Bar) + if l := len(bars); l != 3 { + t.Errorf("Expected len of bars: %d, got: %d", 3, l) + } + + p.Wait() + pq := mpb.PriorityQueue(bars) + + if bar := heap.Pop(&pq).(*mpb.Bar); bar != b { + t.Errorf("Expected bar b, got: %s", identity[bar]) + } + if bar := heap.Pop(&pq).(*mpb.Bar); bar != c { + t.Errorf("Expected bar c, got: %s", identity[bar]) + } + if bar := heap.Pop(&pq).(*mpb.Bar); bar != a { + t.Errorf("Expected bar a, got: %s", identity[bar]) + } +} + +func TestUpdateBarPriorityLazy(t *testing.T) { + shutdown := make(chan interface{}) + refreshCh := make(chan interface{}) + ctx, cancel := context.WithCancel(context.Background()) + p := mpb.NewWithContext(ctx, + mpb.WithOutput(io.Discard), + mpb.WithManualRefresh(refreshCh), + mpb.WithShutdownNotifier(shutdown), + ) + a := p.AddBar(100, mpb.BarPriority(1)) + b := p.AddBar(100, mpb.BarPriority(2)) + c := p.AddBar(100, mpb.BarPriority(3)) + + identity := map[*mpb.Bar]string{ + a: "a", + b: "b", + c: "c", + } + + // UpdateBarPriority with lazy flat needs at least one refresh cycle to + // update order inside underlying binary heap + p.UpdateBarPriority(c, 2, true) + p.UpdateBarPriority(b, 3, true) + + // ommiting following line will fail the test which proves correct lazy behavior + refreshCh <- time.Now() + + go cancel() bars := (<-shutdown).([]*mpb.Bar) if l := len(bars); l != 3 {