diff --git a/starmap_test.go b/starmap_test.go index ab32bc6..7566a24 100644 --- a/starmap_test.go +++ b/starmap_test.go @@ -2,7 +2,9 @@ package starmap import ( "fmt" + "math/rand" "testing" + "time" ) type Miaomiao struct { @@ -39,3 +41,29 @@ func Test_Remote2(t *testing.T) { fmt.Println(client.StoreMap(&Miaomiao{"suki", 1127, false})) fmt.Println(client.GetMap("miaomiao")) } + +func Test_Math(t *testing.T) { + wg := NewWaitGroup(5000) + rand.Seed(time.Now().UnixNano()) + waitfn := func(wg *WaitGroup, num int) { + defer wg.Done() + fmt.Println(num) + time.Sleep(time.Second * time.Duration(2+rand.Intn(3))) + } + for i := 0; i <= 3214670; i++ { + wg.Add(1) + if i == 34567 { + wg.SetMaxWaitNum(50000) + } + if i == 123456 { + wg.SetMaxWaitNum(210456) + } + if i == 323456 { + wg.SetMaxWaitNum(2104562) + } + go waitfn(&wg, i) + } + fmt.Println("Waiting~~~~~~~~~~") + wg.Wait() + fmt.Println(wg.allCount) +} diff --git a/sync.go b/sync.go new file mode 100644 index 0000000..57405b5 --- /dev/null +++ b/sync.go @@ -0,0 +1,55 @@ +package starmap + +import ( + "sync" + "sync/atomic" + "time" +) + +type WaitGroup struct { + wg *sync.WaitGroup + maxCount uint32 + allCount uint32 +} + +func NewWaitGroup(maxCount int) WaitGroup { + return WaitGroup{wg: &sync.WaitGroup{}, maxCount: uint32(maxCount)} +} + +func (swg *WaitGroup) Add(delta int) { + var Udelta uint32 + if delta < 0 { + Udelta = uint32(-delta - 1) + } else { + Udelta = uint32(delta) + } + for { + allC := atomic.LoadUint32(&swg.allCount) + if atomic.LoadUint32(&swg.maxCount) >= allC+uint32(delta) { + if delta < 0 { + atomic.AddUint32(&swg.allCount, ^uint32(Udelta)) + } else { + atomic.AddUint32(&swg.allCount, uint32(Udelta)) + } + break + } + time.Sleep(time.Microsecond) + } + swg.wg.Add(delta) +} + +func (swg *WaitGroup) Done() { + swg.Add(-1) +} + +func (swg *WaitGroup) Wait() { + swg.wg.Wait() +} + +func (swg *WaitGroup) GetMaxWaitNum() int { + return int(atomic.LoadUint32(&swg.maxCount)) +} + +func (swg *WaitGroup) SetMaxWaitNum(num int) { + atomic.AddUint32(&swg.maxCount, uint32(num)) +}