From 0211d609b0ae8b69272cab71e98c95430724a496 Mon Sep 17 00:00:00 2001 From: qydysky <32743305+qydysky@users.noreply.github.com> Date: Sun, 12 Mar 2023 12:13:38 +0800 Subject: [PATCH] fix --- limit/Limit.go | 146 +++++++++++++++++++++++++++++--------------- limit/Limit_test.go | 131 ++++++++++++++++++++++++--------------- 2 files changed, 180 insertions(+), 97 deletions(-) diff --git a/limit/Limit.go b/limit/Limit.go index fcc0e23..5e13218 100644 --- a/limit/Limit.go +++ b/limit/Limit.go @@ -1,79 +1,127 @@ package part import ( - "time" "sync/atomic" + "time" + + signal "github.com/qydysky/part/signal" ) type Limit struct { - maxNum_in_period int - ms_in_period int - ms_to_timeout int - bucket chan struct{} - pre_bucket_token_num int - wait_num int32 + druation time.Duration + druation_timeout time.Duration + bucket chan struct{} + wait_num atomic.Int32 + pre_bucket_token_num atomic.Int64 + maxNum_in_period int + cancel *signal.Signal } // create a Limit Object // it will allow maxNum_in_period requests(call TO()) in ms_in_period. if the request(call TO()) is out of maxNum_in_period,it will wait ms_to_timeout -//ms_to_timeout [>0:will wait ms] [=0:no wait] [<0:will block] -func New(maxNum_in_period,ms_in_period,ms_to_timeout int) (*Limit) { +// ms_to_timeout [>0:will wait ms] [=0:no wait] [<0:will block] +func New(maxNum_in_period int, druation, druation_timeout string) *Limit { object := Limit{ - maxNum_in_period:maxNum_in_period, - ms_in_period:ms_in_period, - ms_to_timeout:ms_to_timeout, - bucket:make(chan struct{},maxNum_in_period), + bucket: make(chan struct{}, maxNum_in_period), + cancel: signal.Init(), + } + + if maxNum_in_period <= 0 { + panic("maxNum_in_period <= 0") + } else { + object.maxNum_in_period = maxNum_in_period + } + + if t, e := time.ParseDuration(druation); e != nil { + panic(e) + } else { + object.druation = t.Abs() } - go func(object *Limit){ - for object.maxNum_in_period > 0 { - for i:=0;i 0; i-- { + select { + case object.bucket <- struct{}{}: + case <-ec: + return + default: + i = 0 + } } } - time.Sleep(time.Duration(object.ms_in_period)*time.Millisecond) - object.pre_bucket_token_num = len(object.bucket) } - close(object.bucket) - }(&object) + }() - //make sure the bucket is full - for object.TK() != maxNum_in_period {} - object.pre_bucket_token_num = len(object.bucket) - return &object } // the func will return true if the request(call TO()) is up to limit and return false if not func (l *Limit) TO() bool { - var AfterMS = func(ReadTimeout int) (c <-chan time.Time) { - if ReadTimeout > 0 { - c = time.NewTimer(time.Millisecond*time.Duration(ReadTimeout)).C - } else if ReadTimeout == 0 { - tc := make(chan time.Time,1) - tc <- time.Now() - close(tc) - c = tc - } - return - } - - atomic.AddInt32(&l.wait_num, 1) - defer atomic.AddInt32(&l.wait_num, -1) - - select { - case <-l.bucket:; - case <-AfterMS(l.ms_to_timeout):return true; + if !l.IsLive() { + return false } + l.wait_num.Add(1) + defer l.wait_num.Add(-1) - return false + if l.druation_timeout < 0 || l.druation == 0 { + <-l.bucket + return false + } else if l.druation_timeout == 0 { + select { + case <-l.bucket: + return false + default: + return true + } + } else { + select { + case <-l.bucket: + return false + case <-time.NewTimer(l.druation_timeout).C: + return true + } + } +} + +func (l *Limit) IsLive() bool { + return l.cancel.Islive() } func (l *Limit) Close() { - l.maxNum_in_period = 0 + l.cancel.Done() } // return the token number of bucket at now @@ -83,9 +131,9 @@ func (l *Limit) TK() int { // return the token number of bucket at previous func (l *Limit) PTK() int { - return l.pre_bucket_token_num + return int(l.pre_bucket_token_num.Load()) } func (l *Limit) WNum() int32 { - return atomic.LoadInt32(&l.wait_num) -} \ No newline at end of file + return l.wait_num.Load() +} diff --git a/limit/Limit_test.go b/limit/Limit_test.go index eff1698..4133325 100644 --- a/limit/Limit_test.go +++ b/limit/Limit_test.go @@ -5,61 +5,96 @@ import ( "time" ) -func Test_1(t *testing.T){ - l := New(10,1000,0) - pass := 0 - for i:=0;i<1500;i+=1{ - go func(){ - if !l.TO() {pass += 1} - }() - time.Sleep(time.Millisecond) - } - if pass!=20 {t.Error(`pass != 20`)} +func Test_6(t *testing.T) { + l := New(2, "1s", "-1s") + t0 := time.Now() + if l.TO() || time.Since(t0) > time.Millisecond { + t.Fatal() + } + if l.TO() || time.Since(t0) > time.Millisecond { + t.Fatal() + } + t0 = time.Now() + if l.TO() || time.Until(t0.Add(time.Second)) > time.Millisecond { + t.Fatal(time.Since(t0)) + } } -func Test_2(t *testing.T){ - l := New(10,1000,1000) - pass := 0 - for i:=0;i<500;i+=1{ - go func(){ - if !l.TO() {pass += 1} - }() - time.Sleep(time.Millisecond) - } - if pass!=10 {t.Error(`pass != 10`,pass)} +func Test_8(t *testing.T) { + l := New(2, "1s", "0s") + t0 := time.Now() + if l.TO() || time.Since(t0) > time.Millisecond { + t.Fatal() + } + if l.TO() || time.Since(t0) > time.Millisecond { + t.Fatal() + } + if !l.TO() { + t.Fatal() + } } -func Test_3(t *testing.T){ - l := New(10,0,0) - pass := 0 - for i:=0;i<500;i+=1{ - go func(){ - if !l.TO() {pass += 1} - }() - time.Sleep(time.Millisecond) - } - t.Log(pass) +func Test_9(t *testing.T) { + l := New(2, "1s", "5ms") + t0 := time.Now() + if l.TO() || time.Since(t0) > time.Millisecond { + t.Fatal() + } + time.Sleep(time.Millisecond * 500) + if l.TO() || time.Until(t0.Add(time.Millisecond*500)) > time.Millisecond { + t.Fatal() + } + time.Sleep(time.Millisecond * 500) + if l.TO() || time.Until(t0.Add(time.Millisecond*505)) > time.Millisecond { + t.Fatal() + } } -func Test_4(t *testing.T){ - l := New(0,0,10) - pass := 0 - for i:=0;i<500;i+=1{ - go func(){ - if !l.TO() {pass += 1} - }() - time.Sleep(time.Millisecond) - } - t.Log(pass) +func Test_1(t *testing.T) { + l := New(2, "0s", "0s") + t0 := time.Now() + if l.TO() || time.Since(t0) > time.Millisecond { + t.Fatal() + } + if l.TO() || time.Until(t0) > time.Millisecond { + t.Fatal() + } + if l.TO() || time.Until(t0) > time.Millisecond { + t.Fatal() + } +} + +func Test_2(t *testing.T) { + l := New(2, "10s", "-1s") + go func() { + time.Sleep(time.Second) + l.Close() + }() + t0 := time.Now() + if l.TO() || time.Since(t0) > time.Millisecond { + t.Fatal() + } + if l.TO() || time.Until(t0) > time.Millisecond { + t.Fatal() + } + if l.TO() || time.Until(t0.Add(time.Second)) > time.Millisecond { + t.Fatal(l.IsLive(), time.Until(t0.Add(time.Second))) + } } -func Test_5(t *testing.T){ - l := New(100,3000,0) - t.Log(l.TK()) - for i:=1;i<=50;i+=1{ +func Test_5(t *testing.T) { + l := New(100, "3s", "0s") + if l.TK() != 100 { + t.Error(`5`, l.TK()) + } + for i := 1; i <= 50; i += 1 { l.TO() } - if l.TK() != 50 {t.Error(`5`,l.TK())} - time.Sleep(time.Second*3) - if l.PTK() != 50 {t.Error(`5`,l.PTK())} -} \ No newline at end of file + if l.TK() != 50 { + t.Error(`5`, l.TK()) + } + time.Sleep(time.Second * 4) + if l.PTK() != 50 { + t.Error(`5`, l.PTK()) + } +} -- 2.39.2