github.com/awesome-flow/flow@v0.0.3-0.20190918184116-508d75d68a2c/pkg/corev1alpha1/actor/throttler.go (about) 1 package actor 2 3 import ( 4 "fmt" 5 "sync" 6 "sync/atomic" 7 _ "unsafe" 8 9 core "github.com/awesome-flow/flow/pkg/corev1alpha1" 10 ) 11 12 //go:noescape 13 //go:linkname nanotime runtime.nanotime 14 func nanotime() int64 15 16 type Throttler struct { 17 name string 18 ctx *core.Context 19 msgkey string 20 msgcost int64 21 bukcap int64 22 timefun func() int64 23 buckets map[string]*int64 24 mutex sync.RWMutex 25 queue chan *core.Message 26 wg sync.WaitGroup 27 } 28 29 var _ core.Actor = (*Throttler)(nil) 30 31 func NewThrottler(name string, ctx *core.Context, params core.Params) (core.Actor, error) { 32 rps, ok := params["rps"] 33 if !ok { 34 return nil, fmt.Errorf("throttler %s is missing `rps` config", name) 35 } 36 37 msgcost := 1000000000 / int64(rps.(int)) 38 bukcap := 1000000000 - msgcost 39 40 t := &Throttler{ 41 name: name, 42 ctx: ctx, 43 msgcost: msgcost, 44 bukcap: bukcap, 45 timefun: nanotime, 46 buckets: map[string]*int64{"": new(int64)}, 47 queue: make(chan *core.Message), 48 } 49 50 if msgkey, ok := params["msgkey"]; ok { 51 t.msgkey = msgkey.(string) 52 } 53 54 return t, nil 55 } 56 57 func (t *Throttler) Name() string { 58 return t.name 59 } 60 61 func (t *Throttler) Start() error { 62 return nil 63 } 64 65 func (t *Throttler) Stop() error { 66 close(t.queue) 67 t.wg.Wait() 68 69 return nil 70 } 71 72 func (t *Throttler) Connect(nthreads int, peer core.Receiver) error { 73 for i := 0; i < nthreads; i++ { 74 t.wg.Add(1) 75 go func() { 76 for msg := range t.queue { 77 if err := peer.Receive(msg); err != nil { 78 msg.Complete(core.MsgStatusFailed) 79 t.ctx.Logger().Error(err.Error()) 80 } 81 } 82 t.wg.Done() 83 }() 84 } 85 86 return nil 87 } 88 89 func (t *Throttler) Receive(msg *core.Message) error { 90 k := "" 91 if len(t.msgkey) > 0 { 92 if v, ok := msg.Meta(t.msgkey); ok { 93 k = v.(string) 94 } 95 } 96 if t.shouldPass(k) { 97 t.queue <- msg 98 return nil 99 } 100 msg.Complete(core.MsgStatusThrottled) 101 102 return nil 103 } 104 105 func (t *Throttler) getBucket(msgkey string) *int64 { 106 t.mutex.RLock() 107 if bucket, ok := t.buckets[msgkey]; ok { 108 t.mutex.RUnlock() 109 return bucket 110 } 111 112 b := new(int64) 113 *b = t.timefun() - 1 114 115 t.mutex.Lock() 116 defer t.mutex.Unlock() 117 118 if bucket, ok := t.buckets[msgkey]; ok { 119 return bucket 120 } 121 t.buckets[msgkey] = b 122 123 return b 124 } 125 126 func (t *Throttler) shouldPass(msgkey string) bool { 127 bucket := t.getBucket(msgkey) 128 129 for l := 0; l < 10; l++ { 130 now := t.timefun() 131 tat := atomic.LoadInt64(bucket) // theoretical arrival time 132 if now < tat-t.bukcap { 133 return false 134 } 135 var newtat int64 136 if now > tat { 137 newtat = now + t.msgcost 138 } else { 139 newtat = tat + t.msgcost 140 } 141 142 if atomic.CompareAndSwapInt64(bucket, tat, newtat) { 143 return true 144 } 145 } 146 147 return false 148 }