From a239f68e54da1d7b4fe3ce216f4a158b05f8684c Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Sun, 30 May 2021 16:52:12 +0800
Subject: [PATCH] add some func

---
 limit/Limit.go | 27 +++++++++++++++++++++++++--
 reqf/Reqf.go   |  4 ++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/limit/Limit.go b/limit/Limit.go
index edc7d09..2907bb6 100644
--- a/limit/Limit.go
+++ b/limit/Limit.go
@@ -2,6 +2,7 @@ package part
 
 import (
 	"time"
+	"sync/atomic"
 )
 
 type Limit struct {
@@ -10,6 +11,7 @@ type Limit struct {
 	ms_to_timeout int
 	bucket chan struct{}
 	pre_bucket_token_num int
+	wait_num int32
 }
 
 // create a Limit Object
@@ -34,6 +36,7 @@ func New(maxNum_in_period,ms_in_period,ms_to_timeout int) (*Limit) {
 			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
@@ -45,13 +48,29 @@ func New(maxNum_in_period,ms_in_period,ms_to_timeout int) (*Limit) {
 
 // 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
+		}
+		return
+	}
+	
+	atomic.AddInt32(&l.wait_num, 1)
+	defer atomic.AddInt32(&l.wait_num, -1)
+	
 	select {
-		case <-l.bucket :;
-		case <-time.After(time.Duration(l.ms_to_timeout)*time.Millisecond):return true;
+		case <-l.bucket:;
+		case <-AfterMS(l.ms_to_timeout):return true;
 	}
+
+
 	return false
 }
 
+func (l *Limit) Close() {
+	l.maxNum_in_period = 0
+}
+
 // return the token number of bucket at now
 func (l *Limit) TK() int {
 	return len(l.bucket)
@@ -60,4 +79,8 @@ func (l *Limit) TK() int {
 // return the token number of bucket at previous
 func (l *Limit) PTK() int {
 	return l.pre_bucket_token_num
+}
+
+func (l *Limit) WNum() int32 {
+	return atomic.LoadInt32(&l.wait_num)
 }
\ No newline at end of file
diff --git a/reqf/Reqf.go b/reqf/Reqf.go
index 3e129bd..a9763d1 100644
--- a/reqf/Reqf.go
+++ b/reqf/Reqf.go
@@ -331,6 +331,10 @@ func IsTimeout(e error) bool {
     return errors.Is(e, context.DeadlineExceeded)
 }
 
+func IsDnsErr(e error) bool {
+    return e != nil && strings.Contains(e.Error(), "lookup")
+}
+
 func IsCancel(e error) bool {
     return errors.Is(e, context.Canceled)
 }
\ No newline at end of file
-- 
2.39.2