exitgroup.go

package main

import "sync"

type ExitGroup struct {
	wg sync.WaitGroup
	cl []func()
	lk sync.Mutex
}

func (e *ExitGroup) OnClose(fn func()) {
	e.lk.Lock()
	e.wg.Add(1)
	e.cl = append(e.cl, fn)
	e.lk.Unlock()
}

func (e *ExitGroup) OnCloseErr(fn func() error) {
	e.OnClose(func() {
		if err := fn(); err != nil {
			log.Println(err)
		}
	})
}

func (e *ExitGroup) Exit() {
	e.lk.Lock()
	for _, fn := range e.cl {
		fn()
		e.wg.Done()
	}
	e.lk.Unlock()
}

func (e *ExitGroup) Wait() {
	e.wg.Wait()
}