initial import
This commit is contained in:
kevin
446
rq/pusher.go
Normal file
446
rq/pusher.go
Normal file
@@ -0,0 +1,446 @@
|
||||
package rq
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"zero/core/discov"
|
||||
"zero/core/errorx"
|
||||
"zero/core/jsonx"
|
||||
"zero/core/lang"
|
||||
"zero/core/logx"
|
||||
"zero/core/queue"
|
||||
"zero/core/redisqueue"
|
||||
"zero/core/stores/redis"
|
||||
"zero/core/threading"
|
||||
"zero/rq/constant"
|
||||
"zero/rq/update"
|
||||
)
|
||||
|
||||
const (
|
||||
retryTimes = 3
|
||||
etcdRedisFields = 4
|
||||
)
|
||||
|
||||
var ErrPusherTypeError = errors.New("not a QueuePusher instance")
|
||||
|
||||
type (
|
||||
KeyFn func(string) (key, payload string, err error)
|
||||
KeysFn func(string) (ctx context.Context, keys []string, err error)
|
||||
AssembleFn func(context.Context, []string) (payload string, err error)
|
||||
PusherOption func(*Pusher) error
|
||||
|
||||
// just push once or do it retryTimes, it's a choice.
|
||||
// because only when at least a server is alive, and
|
||||
// pushing to the server failed, we'll return with an error
|
||||
// if waken up, but the server is going down very quickly,
|
||||
// we're going to wait again. so it's safe to push once.
|
||||
pushStrategy interface {
|
||||
addListener(listener discov.Listener)
|
||||
push(string) error
|
||||
}
|
||||
|
||||
batchConsistentStrategy struct {
|
||||
keysFn KeysFn
|
||||
assembleFn AssembleFn
|
||||
subClient *discov.BatchConsistentSubClient
|
||||
}
|
||||
|
||||
consistentStrategy struct {
|
||||
keyFn KeyFn
|
||||
subClient *discov.ConsistentSubClient
|
||||
}
|
||||
|
||||
roundRobinStrategy struct {
|
||||
subClient *discov.RoundRobinSubClient
|
||||
}
|
||||
|
||||
serverListener struct {
|
||||
updater *update.IncrementalUpdater
|
||||
}
|
||||
|
||||
Pusher struct {
|
||||
name string
|
||||
endpoints []string
|
||||
key string
|
||||
failovers sync.Map
|
||||
strategy pushStrategy
|
||||
serverSensitive bool
|
||||
}
|
||||
)
|
||||
|
||||
func NewPusher(endpoints []string, key string, opts ...PusherOption) (*Pusher, error) {
|
||||
pusher := &Pusher{
|
||||
name: getName(key),
|
||||
endpoints: endpoints,
|
||||
key: key,
|
||||
}
|
||||
|
||||
if len(opts) == 0 {
|
||||
opts = []PusherOption{WithRoundRobinStrategy()}
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
if err := opt(pusher); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
if pusher.serverSensitive {
|
||||
listener := new(serverListener)
|
||||
listener.updater = update.NewIncrementalUpdater(listener.update)
|
||||
pusher.strategy.addListener(listener)
|
||||
}
|
||||
|
||||
return pusher, nil
|
||||
}
|
||||
|
||||
func (pusher *Pusher) Name() string {
|
||||
return pusher.name
|
||||
}
|
||||
|
||||
func (pusher *Pusher) Push(message string) error {
|
||||
return pusher.strategy.push(message)
|
||||
}
|
||||
|
||||
func (pusher *Pusher) close(server string, conn interface{}) error {
|
||||
logx.Errorf("dropped redis node: %s", server)
|
||||
|
||||
return pusher.failover(server)
|
||||
}
|
||||
|
||||
func (pusher *Pusher) dial(server string) (interface{}, error) {
|
||||
pusher.failovers.Delete(server)
|
||||
|
||||
p, err := newPusher(server)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
logx.Infof("new redis node: %s", server)
|
||||
|
||||
return p, nil
|
||||
}
|
||||
|
||||
func (pusher *Pusher) failover(server string) error {
|
||||
pusher.failovers.Store(server, lang.Placeholder)
|
||||
|
||||
rds, key, option, err := newRedisWithKey(server)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
threading.GoSafe(func() {
|
||||
defer pusher.failovers.Delete(server)
|
||||
|
||||
for {
|
||||
_, ok := pusher.failovers.Load(server)
|
||||
if !ok {
|
||||
logx.Infof("redis queue (%s) revived", server)
|
||||
return
|
||||
}
|
||||
|
||||
message, err := rds.Lpop(key)
|
||||
if err != nil {
|
||||
logx.Error(err)
|
||||
return
|
||||
}
|
||||
|
||||
if len(message) == 0 {
|
||||
logx.Infof("repush redis queue (%s) done", server)
|
||||
return
|
||||
}
|
||||
|
||||
if option == constant.TimedQueueType {
|
||||
message, err = unwrapTimedMessage(message)
|
||||
if err != nil {
|
||||
logx.Errorf("invalid timedMessage: %s, error: %s", message, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if err = pusher.strategy.push(message); err != nil {
|
||||
logx.Error(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func UnmarshalPusher(server string) (queue.QueuePusher, error) {
|
||||
store, key, option, err := newRedisWithKey(server)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if option == constant.TimedQueueType {
|
||||
return redisqueue.NewPusher(store, key, redisqueue.WithTime()), nil
|
||||
}
|
||||
|
||||
return redisqueue.NewPusher(store, key), nil
|
||||
}
|
||||
|
||||
func WithBatchConsistentStrategy(keysFn KeysFn, assembleFn AssembleFn, opts ...discov.BalanceOption) PusherOption {
|
||||
return func(pusher *Pusher) error {
|
||||
subClient, err := discov.NewBatchConsistentSubClient(pusher.endpoints, pusher.key, pusher.dial,
|
||||
pusher.close, opts...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
pusher.strategy = batchConsistentStrategy{
|
||||
keysFn: keysFn,
|
||||
assembleFn: assembleFn,
|
||||
subClient: subClient,
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func WithConsistentStrategy(keyFn KeyFn, opts ...discov.BalanceOption) PusherOption {
|
||||
return func(pusher *Pusher) error {
|
||||
subClient, err := discov.NewConsistentSubClient(pusher.endpoints, pusher.key, pusher.dial, pusher.close, opts...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
pusher.strategy = consistentStrategy{
|
||||
keyFn: keyFn,
|
||||
subClient: subClient,
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func WithRoundRobinStrategy() PusherOption {
|
||||
return func(pusher *Pusher) error {
|
||||
subClient, err := discov.NewRoundRobinSubClient(pusher.endpoints, pusher.key, pusher.dial, pusher.close)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
pusher.strategy = roundRobinStrategy{
|
||||
subClient: subClient,
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func WithServerSensitive() PusherOption {
|
||||
return func(pusher *Pusher) error {
|
||||
pusher.serverSensitive = true
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (bcs batchConsistentStrategy) addListener(listener discov.Listener) {
|
||||
bcs.subClient.AddListener(listener)
|
||||
}
|
||||
|
||||
func (bcs batchConsistentStrategy) balance(keys []string) map[interface{}][]string {
|
||||
// we need to make sure the servers are available, otherwise wait forever
|
||||
for {
|
||||
if mapping, ok := bcs.subClient.Next(keys); ok {
|
||||
return mapping
|
||||
} else {
|
||||
bcs.subClient.WaitForServers()
|
||||
// make sure we don't flood logs too much in extreme conditions
|
||||
time.Sleep(time.Second)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (bcs batchConsistentStrategy) push(message string) error {
|
||||
ctx, keys, err := bcs.keysFn(message)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var batchError errorx.BatchError
|
||||
mapping := bcs.balance(keys)
|
||||
for conn, connKeys := range mapping {
|
||||
payload, err := bcs.assembleFn(ctx, connKeys)
|
||||
if err != nil {
|
||||
batchError.Add(err)
|
||||
continue
|
||||
}
|
||||
|
||||
for i := 0; i < retryTimes; i++ {
|
||||
if err = bcs.pushOnce(conn, payload); err != nil {
|
||||
batchError.Add(err)
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return batchError.Err()
|
||||
}
|
||||
|
||||
func (bcs batchConsistentStrategy) pushOnce(server interface{}, payload string) error {
|
||||
pusher, ok := server.(queue.QueuePusher)
|
||||
if ok {
|
||||
return pusher.Push(payload)
|
||||
} else {
|
||||
return ErrPusherTypeError
|
||||
}
|
||||
}
|
||||
|
||||
func (cs consistentStrategy) addListener(listener discov.Listener) {
|
||||
cs.subClient.AddListener(listener)
|
||||
}
|
||||
|
||||
func (cs consistentStrategy) push(message string) error {
|
||||
var batchError errorx.BatchError
|
||||
|
||||
key, payload, err := cs.keyFn(message)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for i := 0; i < retryTimes; i++ {
|
||||
if err = cs.pushOnce(key, payload); err != nil {
|
||||
batchError.Add(err)
|
||||
} else {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return batchError.Err()
|
||||
}
|
||||
|
||||
func (cs consistentStrategy) pushOnce(key, payload string) error {
|
||||
// we need to make sure the servers are available, otherwise wait forever
|
||||
for {
|
||||
if server, ok := cs.subClient.Next(key); ok {
|
||||
pusher, ok := server.(queue.QueuePusher)
|
||||
if ok {
|
||||
return pusher.Push(payload)
|
||||
} else {
|
||||
return ErrPusherTypeError
|
||||
}
|
||||
} else {
|
||||
cs.subClient.WaitForServers()
|
||||
// make sure we don't flood logs too much in extreme conditions
|
||||
time.Sleep(time.Second)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (rrs roundRobinStrategy) addListener(listener discov.Listener) {
|
||||
rrs.subClient.AddListener(listener)
|
||||
}
|
||||
|
||||
func (rrs roundRobinStrategy) push(message string) error {
|
||||
var batchError errorx.BatchError
|
||||
|
||||
for i := 0; i < retryTimes; i++ {
|
||||
if err := rrs.pushOnce(message); err != nil {
|
||||
batchError.Add(err)
|
||||
} else {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return batchError.Err()
|
||||
}
|
||||
|
||||
func (rrs roundRobinStrategy) pushOnce(message string) error {
|
||||
if server, ok := rrs.subClient.Next(); ok {
|
||||
pusher, ok := server.(queue.QueuePusher)
|
||||
if ok {
|
||||
return pusher.Push(message)
|
||||
} else {
|
||||
return ErrPusherTypeError
|
||||
}
|
||||
} else {
|
||||
rrs.subClient.WaitForServers()
|
||||
return rrs.pushOnce(message)
|
||||
}
|
||||
}
|
||||
|
||||
func getName(key string) string {
|
||||
return fmt.Sprintf("etcd:%s", key)
|
||||
}
|
||||
|
||||
func newPusher(server string) (queue.QueuePusher, error) {
|
||||
if rds, key, option, err := newRedisWithKey(server); err != nil {
|
||||
return nil, err
|
||||
} else if option == constant.TimedQueueType {
|
||||
return redisqueue.NewPusher(rds, key, redisqueue.WithTime()), nil
|
||||
} else {
|
||||
return redisqueue.NewPusher(rds, key), nil
|
||||
}
|
||||
}
|
||||
|
||||
func newRedisWithKey(server string) (rds *redis.Redis, key, option string, err error) {
|
||||
fields := strings.Split(server, constant.Delimeter)
|
||||
if len(fields) < etcdRedisFields {
|
||||
err = fmt.Errorf("wrong redis queue: %s, should be ip:port/type/password/key/[option]", server)
|
||||
return
|
||||
}
|
||||
|
||||
addr := fields[0]
|
||||
tp := fields[1]
|
||||
pass := fields[2]
|
||||
key = fields[3]
|
||||
|
||||
if len(fields) > etcdRedisFields {
|
||||
option = fields[4]
|
||||
}
|
||||
|
||||
rds = redis.NewRedis(addr, tp, pass)
|
||||
return
|
||||
}
|
||||
|
||||
func (sl *serverListener) OnUpdate(keys []string, servers []string, newKey string) {
|
||||
sl.updater.Update(keys, servers, newKey)
|
||||
}
|
||||
|
||||
func (sl *serverListener) OnReload() {
|
||||
sl.updater.Update(nil, nil, "")
|
||||
}
|
||||
|
||||
func (sl *serverListener) update(change update.ServerChange) {
|
||||
content, err := change.Marshal()
|
||||
if err != nil {
|
||||
logx.Error(err)
|
||||
}
|
||||
|
||||
if err = broadcast(change.Servers, content); err != nil {
|
||||
logx.Error(err)
|
||||
}
|
||||
}
|
||||
|
||||
func broadcast(servers []string, message string) error {
|
||||
var be errorx.BatchError
|
||||
|
||||
for _, server := range servers {
|
||||
q, err := UnmarshalPusher(server)
|
||||
if err != nil {
|
||||
be.Add(err)
|
||||
} else {
|
||||
q.Push(message)
|
||||
}
|
||||
}
|
||||
|
||||
return be.Err()
|
||||
}
|
||||
|
||||
func unwrapTimedMessage(message string) (string, error) {
|
||||
var tm redisqueue.TimedMessage
|
||||
if err := jsonx.UnmarshalFromString(message, &tm); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return tm.Payload, nil
|
||||
}
|
||||
Reference in New Issue
Block a user