diff --git a/bar.go b/bar.go index ac1b9c1..0428f71 100644 --- a/bar.go +++ b/bar.go @@ -31,27 +31,26 @@ // bState is actual bar's state. type bState struct { - id int - priority int - reqWidth int - shutdown int - total int64 - current int64 - refill int64 - trimSpace bool - aborted bool - triggerComplete bool - rmOnComplete bool - noPop bool - autoRefresh bool - buffers [3]*bytes.Buffer - decorators [2][]decor.Decorator - ewmaDecorators []decor.EwmaDecorator - shutdownListeners []decor.ShutdownListener - filler BarFiller - extender extenderFunc - renderReq chan<- time.Time - waitBar *Bar // key for (*pState).queueBars + id int + priority int + reqWidth int + shutdown int + total int64 + current int64 + refill int64 + trimSpace bool + aborted bool + triggerComplete bool + rmOnComplete bool + noPop bool + autoRefresh bool + buffers [3]*bytes.Buffer + decorators [2][]decor.Decorator + ewmaDecorators []decor.EwmaDecorator + filler BarFiller + extender extenderFunc + renderReq chan<- time.Time + waitBar *Bar // key for (*pState).queueBars } type renderFrame struct { @@ -401,19 +400,24 @@ } func (b *Bar) serve(bs *bState) { + decoratorsOnShutdown := func(decorators []decor.Decorator) { + for _, d := range decorators { + if d, ok := unwrap(d).(decor.ShutdownListener); ok { + b.container.bwg.Add(1) + go func() { + d.OnShutdown() + b.container.bwg.Done() + }() + } + } + } for { select { case op := <-b.operateState: op(bs) case <-b.ctx.Done(): - for _, d := range bs.shutdownListeners { - b.container.bwg.Add(1) - d := d - go func() { - d.OnShutdown() - b.container.bwg.Done() - }() - } + decoratorsOnShutdown(bs.decorators[0]) + decoratorsOnShutdown(bs.decorators[1]) bs.aborted = !bs.completed() b.bs = bs close(b.bsOk) @@ -564,9 +568,6 @@ if d, ok := d.(decor.EwmaDecorator); ok { s.ewmaDecorators = append(s.ewmaDecorators, d) } - if d, ok := d.(decor.ShutdownListener); ok { - s.shutdownListeners = append(s.shutdownListeners, d) - } } }