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  }