GoDoxy/internal/task/with.go

48 lines
701 B
Go

package task
import (
"sync"
"github.com/puzpuzpuz/xsync/v4"
)
type withWg[T comparable] struct {
m *xsync.Map[T, struct{}]
wg sync.WaitGroup
}
func newWithWg[T comparable]() *withWg[T] {
return &withWg[T]{
m: xsync.NewMap[T, struct{}](),
}
}
func (w *withWg[T]) Add(ele T) {
w.wg.Add(1)
w.m.Store(ele, struct{}{})
}
func (w *withWg[T]) AddWithoutWG(ele T) {
w.m.Store(ele, struct{}{})
}
func (w *withWg[T]) Delete(key T) {
w.wg.Done()
w.m.Delete(key)
}
func (w *withWg[T]) DeleteWithoutWG(key T) {
w.m.Delete(key)
}
func (w *withWg[T]) Wait() {
w.wg.Wait()
}
func (w *withWg[T]) Range(yield func(T) bool) {
for ele := range w.m.Range {
if !yield(ele) {
break
}
}
}