0g-chain/helper/da/utils/sizedw8grp/sizedw8grp.go
2024-08-02 19:26:50 +08:00

52 lines
735 B
Go

package sizedw8grp
import (
"context"
"math"
"sync"
)
type SizedWaitGroup struct {
Size int
current chan struct{}
wg sync.WaitGroup
}
func New(limit int) SizedWaitGroup {
size := math.MaxInt32
if limit > 0 {
size = limit
}
return SizedWaitGroup{
Size: size,
current: make(chan struct{}, size),
wg: sync.WaitGroup{},
}
}
func (s *SizedWaitGroup) Add() {
_ = s.AddWithContext(context.Background())
}
func (s *SizedWaitGroup) AddWithContext(ctx context.Context) error {
select {
case <-ctx.Done():
return ctx.Err()
case s.current <- struct{}{}:
break
}
s.wg.Add(1)
return nil
}
func (s *SizedWaitGroup) Done() {
<-s.current
s.wg.Done()
}
func (s *SizedWaitGroup) Wait() {
s.wg.Wait()
}