add mtx

Browse Source
main
git 2024-06-10 18:39:47 +08:00
parent 8ec1227f87
commit d8b9d33b82
Signed by: git
GPG Key ID: 3F65EFFA44207ADD
1 changed files with 25 additions and 11 deletions

View File

@ -116,17 +116,15 @@ func (b *etcdBroker) StartConsuming(consumerTag string, concurrency int, taskPro
continue
}
task, err := b.nextTask(getQueue(b.GetConfig(), taskProcessor), consumerTag)
if err != nil {
log.ERROR.Printf("get next task failed: %s", err)
task := b.nextTask(getQueue(b.GetConfig(), taskProcessor), consumerTag)
if task == nil {
time.Sleep(time.Second)
continue
}
if task != nil {
deliveries <- task
}
}
}
}()
// A goroutine to watch for delayed tasks and push them to deliveries
@ -297,8 +295,15 @@ func (b *etcdBroker) GetDelayedTasks() ([]*tasks.Signature, error) {
return items, nil
}
func (b *etcdBroker) nextTask(queue string, consumerTag string) (Delivery, error) {
for k, assigned := range b.assignMap {
func (b *etcdBroker) nextTask(queue string, consumerTag string) Delivery {
b.mtx.Lock()
assignMap := make(map[string]bool, len(b.assignMap))
for k, v := range assignMap {
assignMap[k] = v
}
b.mtx.Unlock()
for k, assigned := range assignMap {
if assigned {
continue
}
@ -311,11 +316,10 @@ func (b *etcdBroker) nextTask(queue string, consumerTag string) (Delivery, error
continue
}
return d, nil
return d
}
time.Sleep(time.Second)
return b.nextTask(queue, consumerTag)
return nil
}
func (b *etcdBroker) setAssign(key string, assign bool) bool {
@ -350,7 +354,9 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
if b.setAssign(key, true) {
continue
}
b.mtx.Lock()
b.assignMap[key] = false
b.mtx.Unlock()
}
// Watch
@ -358,6 +364,10 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
defer watchCancel()
wc := b.client.Watch(watchCtx, keyPrefix, clientv3.WithPrefix(), clientv3.WithKeysOnly(), clientv3.WithRev(resp.Header.Revision))
for wresp := range wc {
if wresp.Err() != nil {
return watchCtx.Err()
}
for _, ev := range wresp.Events {
key := string(ev.Kv.Key)
if ev.Type == clientv3.EventTypeDelete {
@ -365,7 +375,9 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
continue
}
b.mtx.Lock()
delete(b.assignMap, key)
b.mtx.Unlock()
}
if ev.Type == clientv3.EventTypePut {
@ -373,7 +385,9 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
continue
}
b.mtx.Lock()
b.assignMap[key] = false
b.mtx.Unlock()
}
}
}