github.com/lingyao2333/mo-zero@v1.4.1/core/queue/balancedpusher.go (about)

     1  package queue
     2  
     3  import (
     4  	"errors"
     5  	"sync/atomic"
     6  
     7  	"github.com/lingyao2333/mo-zero/core/logx"
     8  )
     9  
    10  // ErrNoAvailablePusher indicates no pusher available.
    11  var ErrNoAvailablePusher = errors.New("no available pusher")
    12  
    13  // A BalancedPusher is used to push messages to multiple pusher with round robin algorithm.
    14  type BalancedPusher struct {
    15  	name    string
    16  	pushers []Pusher
    17  	index   uint64
    18  }
    19  
    20  // NewBalancedPusher returns a BalancedPusher.
    21  func NewBalancedPusher(pushers []Pusher) Pusher {
    22  	return &BalancedPusher{
    23  		name:    generateName(pushers),
    24  		pushers: pushers,
    25  	}
    26  }
    27  
    28  // Name returns the name of pusher.
    29  func (pusher *BalancedPusher) Name() string {
    30  	return pusher.name
    31  }
    32  
    33  // Push pushes message to one of the underlying pushers.
    34  func (pusher *BalancedPusher) Push(message string) error {
    35  	size := len(pusher.pushers)
    36  
    37  	for i := 0; i < size; i++ {
    38  		index := atomic.AddUint64(&pusher.index, 1) % uint64(size)
    39  		target := pusher.pushers[index]
    40  
    41  		if err := target.Push(message); err != nil {
    42  			logx.Error(err)
    43  		} else {
    44  			return nil
    45  		}
    46  	}
    47  
    48  	return ErrNoAvailablePusher
    49  }