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<object.maxNum_in_period;i++{
+ if t, e := time.ParseDuration(druation_timeout); e != nil {
+ panic(e)
+ } else {
+ object.druation_timeout = t
+ }
+
+ for i := 0; i < maxNum_in_period; i++ {
+ object.bucket <- struct{}{}
+ }
+
+ go func() {
+ ec, fin := object.cancel.WaitC()
+ defer fin()
+ defer close(object.bucket)
+
+ if object.druation == 0 {
+ for {
+ select {
+ case object.bucket <- struct{}{}:
+ case <-ec:
+ return
+ }
+ }
+ } else {
+ for {
select {
- case object.bucket <- struct{}{}:;
- default :i = object.maxNum_in_period
+ case <-time.After(object.druation):
+ case <-ec:
+ return
+ }
+
+ object.pre_bucket_token_num.Store(int64(len(object.bucket)))
+
+ for i := object.maxNum_in_period; 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
// 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()
+}
"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())
+ }
+}
projects / part / .git / commitdiff