From 7aba7acfe3633eb02c0ca47d8af1d66e7c4175d4 Mon Sep 17 00:00:00 2001
From: zenghongru <zenghongru@gmail.com>
Date: Thu, 27 Jun 2019 14:07:26 +0800
Subject: [PATCH] fix: long pull rate

---
 api.go       |  2 +-
 diamond.go   | 33 +++++++++++++++++++++++----------
 long_pull.go |  3 ++-
 setter.go    | 12 ++++++++++++
 4 files changed, 38 insertions(+), 12 deletions(-)
 create mode 100644 setter.go

diff --git a/api.go b/api.go
index 3bb52da..8385f15 100644
--- a/api.go
+++ b/api.go
@@ -13,7 +13,7 @@ const (
 	acmDeleteConfig  URL = "/diamond-server/datum.do?method=deleteAllDatums"
 	acmPublishConfig URL = "/diamond-server/basestone.do?method=syncUpdateAll"
 	acmAllConfig     URL = "/diamond-server/basestone.do?method=getAllConfigByTenant"
-	acmLongPoll      URL = "/diamond-server/config.co"
+	acmLongPull      URL = "/diamond-server/config.co"
 )
 
 func (u URL) String(addr string) string {
diff --git a/diamond.go b/diamond.go
index bff911c..6d137fd 100644
--- a/diamond.go
+++ b/diamond.go
@@ -4,6 +4,8 @@ import (
 	"math/rand"
 	"time"
 
+	"go.uber.org/ratelimit"
+
 	"github.com/sirupsen/logrus"
 
 	"github.com/xiaojiaoyu100/cast"
@@ -60,15 +62,17 @@ type Config struct {
 
 // Diamond 提供了操作阿里云ACM的能力
 type Diamond struct {
-	option  Option
-	c       *cast.Cast
-	units   []Unit
-	errHook Hook
-	r       *rand.Rand
+	option              Option
+	c                   *cast.Cast
+	units               []Unit
+	errHook             Hook
+	r                   *rand.Rand
+	longPullRate        int
+	longPullRateLimiter ratelimit.Limiter
 }
 
 // New 产生Diamond实例
-func New(addr, tenant, accessKey, secretKey string) (*Diamond, error) {
+func New(addr, tenant, accessKey, secretKey string, setters ...Setter) (*Diamond, error) {
 	option := Option{
 		addr:      addr,
 		tenant:    tenant,
@@ -87,15 +91,24 @@ func New(addr, tenant, accessKey, secretKey string) (*Diamond, error) {
 	if err != nil {
 		return nil, err
 	}
-
 	s := rand.NewSource(time.Now().UnixNano())
 	r := rand.New(s)
 
 	d := &Diamond{
-		option: option,
-		c:      c,
-		r:      r,
+		option:       option,
+		c:            c,
+		r:            r,
+		longPullRate: 6,
 	}
+
+	for _, setter := range setters {
+		if err := setter(d); err != nil {
+			return nil, err
+		}
+	}
+
+	d.longPullRateLimiter = ratelimit.New(d.longPullRate)
+
 	return d, nil
 }
 
diff --git a/long_pull.go b/long_pull.go
index 97ee783..3cb776b 100644
--- a/long_pull.go
+++ b/long_pull.go
@@ -14,6 +14,7 @@ const (
 
 // LongPull 监听配置
 func (d *Diamond) LongPull(unit Unit, contentMD5 string) (string, error) {
+	d.longPullRateLimiter.Take()
 	ip, err := d.QueryIP()
 	if err != nil {
 		return "", err
@@ -33,7 +34,7 @@ func (d *Diamond) LongPull(unit Unit, contentMD5 string) (string, error) {
 	}
 	longPollRequest.ProbeModifyRequest = strings.Join([]string{unit.DataID, unit.Group, contentMD5, d.option.tenant}, wordSeparator) + lineSeparator
 	request := d.c.NewRequest().
-		WithPath(acmLongPoll.String(ip)).
+		WithPath(acmLongPull.String(ip)).
 		WithFormURLEncodedBody(longPollRequest).
 		WithHeader(header).
 		Post()
diff --git a/setter.go b/setter.go
new file mode 100644
index 0000000..d4a076b
--- /dev/null
+++ b/setter.go
@@ -0,0 +1,12 @@
+package aliacm
+
+// Setter configures the diamond.
+type Setter func(d *Diamond) error
+
+// WithLongPullRate sets long pull rate.
+func WithLongPullRate(rate int) Setter {
+	return func(d *Diamond) error {
+		d.longPullRate = rate
+		return nil
+	}
+}