common
/
machinery-plugins
Private
1
Fork 0
Code Pull Requests Releases Activity

finish handleDelayedTask

Browse Source
Operations
Operations
Create branch
Create tag
Revert
Cherry-pick
main
git 2024-06-08 09:47:54 +08:00
parent daa0b5a0ba
commit 3b550dd8e5
Signed by: git
GPG Key ID: 3F65EFFA44207ADD
1 changed files with 96 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

117
brokers/etcd/etcd.go
View File

@ -30,6 +30,8 @@ type etcdBroker struct {
ctx context.Context
client *clientv3.Client
wg sync.WaitGroup
keyMap map[string]struct{}
mtx sync.Mutex
}
// New ..
@ -50,6 +52,7 @@ func New(ctx context.Context, conf *config.Config) (iface.Broker, error) {
Broker: common.NewBroker(conf),
ctx: ctx,
client: client,
keyMap: map[string]struct{}{},
}
return &broker, nil
@ -68,12 +71,36 @@ func (b *etcdBroker) StartConsuming(consumerTag string, concurrency int, taskPro
// Channel to which we will push tasks ready for processing by worker
deliveries := make(chan Delivery)
ctx, cancel := context.WithCancel(b.ctx)
defer cancel()
// list watch task
b.wg.Add(1)
go func() {
defer b.wg.Done()
for {
select {
// A way to stop this goroutine from b.StopConsuming
case <-b.GetStopChan():
return
default:
err := b.listWatchTasks(ctx, getQueue(b.GetConfig(), taskProcessor))
if err != nil {
log.ERROR.Printf("handle list watch task err: %s", err)
}
}
}
}()
// A receiving goroutine keeps popping messages from the queue by BLPOP
// If the message is valid and can be unmarshaled into a proper structure
// we send it to the deliveries channel
b.wg.Add(1)
go func() {
defer b.wg.Done()
defer cancel()
for {
select {
@ -107,6 +134,8 @@ func (b *etcdBroker) StartConsuming(consumerTag string, concurrency int, taskPro
b.wg.Add(1)
go func() {
defer b.wg.Done()
defer cancel()
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
@ -117,7 +146,7 @@ func (b *etcdBroker) StartConsuming(consumerTag string, concurrency int, taskPro
return
case <-ticker.C:
err := b.handleDelayedTask()
err := b.handleDelayedTask(ctx)
if err != nil {
log.ERROR.Printf("handleDelayedTask err: %s", err)
}
@ -208,7 +237,8 @@ func (b *etcdBroker) Publish(ctx context.Context, signature *tasks.Signature) er
// Check the ETA signature field, if it is set and it is in the future,
// delay the task
if signature.ETA != nil && signature.ETA.After(now) {
key = fmt.Sprintf("/machinery/v2/broker/delayed_tasks/eta-%d/%s/%s", signature.ETA.UnixMilli(), signature.RoutingKey, signature.UUID)
key = fmt.Sprintf("/machinery/v2/broker/delayed_tasks/eta-%d/%s/%s",
signature.ETA.UnixMilli(), signature.RoutingKey, signature.UUID)
_, err = b.client.Put(ctx, key, string(msg))
return err
}
@ -245,7 +275,7 @@ func (b *etcdBroker) GetPendingTasks(queue string) ([]*tasks.Signature, error) {
}
key := fmt.Sprintf("/machinery/v2/broker/pending_tasks/%s", queue)
items, err := b.getTasks(context.Background(), key)
items, err := b.getTasks(b.ctx, key)
if err != nil {
return nil, err
}
@ -257,7 +287,7 @@ func (b *etcdBroker) GetPendingTasks(queue string) ([]*tasks.Signature, error) {
func (b *etcdBroker) GetDelayedTasks() ([]*tasks.Signature, error) {
key := "/machinery/v2/broker/delayed_tasks"
items, err := b.getTasks(context.Background(), key)
items, err := b.getTasks(b.ctx, key)
if err != nil {
return nil, err
}
@ -266,33 +296,72 @@ func (b *etcdBroker) GetDelayedTasks() ([]*tasks.Signature, error) {
}
func (b *etcdBroker) nextTask(queue string) (Delivery, error) {
keyPrefix := fmt.Sprintf("/machinery/v2/broker/pending_tasks/%s", queue)
ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
defer cancel()
item, err := getFirstItem(ctx, b.client, keyPrefix)
if err != nil {
return nil, err
for k := range b.keyMap {
if !strings.Contains(k, queue) {
continue
}
signature := new(tasks.Signature)
decoder := json.NewDecoder(bytes.NewReader(item))
decoder.UseNumber()
if err := decoder.Decode(signature); err != nil {
return nil, errs.NewErrCouldNotUnmarshalTaskSignature(item, err)
}
d, err := NewDelivery(b.ctx, b.client, "")
d, err := NewDelivery(b.ctx, b.client, k)
if err != nil {
return nil, err
continue
}
return d, nil
}
func (b *etcdBroker) handleDelayedTask() error {
time.Sleep(time.Second)
return b.nextTask(queue)
}
func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
keyPrefix := fmt.Sprintf("/machinery/v2/broker/pending_tasks/%s", queue)
// List
listCtx, listCancel := context.WithTimeout(ctx, time.Second*5)
defer listCancel()
resp, err := b.client.Get(listCtx, keyPrefix, clientv3.WithPrefix(), clientv3.WithKeysOnly())
if err != nil {
return err
}
for _, kv := range resp.Kvs {
if bytes.Contains(kv.Key, []byte("/assign")) {
delete(b.keyMap, string(kv.Key)[:23])
continue
}
b.keyMap[string(kv.Key)] = struct{}{}
}
// Watch
watchCtx, watchCancel := context.WithTimeout(ctx, time.Minute*60)
defer watchCancel()
wc := b.client.Watch(watchCtx, keyPrefix, clientv3.WithPrefix(), clientv3.WithKeysOnly(), clientv3.WithRev(resp.Header.Revision))
for wresp := range wc {
for _, ev := range wresp.Events {
if ev.Type == clientv3.EventTypeDelete {
if bytes.Contains(ev.Kv.Key, []byte("/assign")) {
b.keyMap[string(ev.Kv.Key)] = struct{}{}
}
continue
}
if ev.Type == clientv3.EventTypePut {
if bytes.Contains(ev.Kv.Key, []byte("/assign")) {
delete(b.keyMap, string(ev.Kv.Key)[:23])
continue
}
b.keyMap[string(ev.Kv.Key)] = struct{}{}
}
}
}
return nil
}
func (b *etcdBroker) handleDelayedTask(ctx context.Context) error {
ttl := time.Second * 10
ctx, cancel := context.WithTimeout(b.ctx, ttl)
ctx, cancel := context.WithTimeout(ctx, ttl)
defer cancel()
// 创建一个新的session
@ -321,13 +390,13 @@ func (b *etcdBroker) handleDelayedTask() error {
for _, kv := range kvs.Kvs {
key := string(kv.Key)
parts := strings.Split(key, "/")
if len(parts) != 7 {
if len(parts) != 8 {
log.WARNING.Printf("invalid delay task %s, continue", key)
continue
}
cmp := clientv3.Compare(clientv3.ModRevision(key), "=", kv.ModRevision)
deleteReq := clientv3.OpDelete(key)
pendingKey := fmt.Sprintf("/machinery/v2/broker/pending_tasks/%s/%s", parts[5], parts[6])
pendingKey := fmt.Sprintf("/machinery/v2/broker/pending_tasks/%s/%s", parts[6], parts[7])
putReq := clientv3.OpPut(pendingKey, string(kv.Value))
c, err := b.client.Txn(b.ctx).If(cmp).Then(deleteReq, putReq).Commit()
if err != nil {