github.com/MeteorsLiu/simpleMQ@v1.0.3/queue/poll.go (about)

     1  package queue
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/alphadose/haxmap"
     7  )
     8  
     9  var (
    10  	ErrPollableEventsNotExists = fmt.Errorf("the pollable event doesn't exist")
    11  )
    12  
    13  type Pollable interface {
    14  	Task
    15  }
    16  
    17  type PollTask struct {
    18  	m *haxmap.Map[string, Pollable]
    19  }
    20  
    21  func NewPoll() *PollTask {
    22  	return &PollTask{
    23  		m: haxmap.New[string, Pollable](),
    24  	}
    25  }
    26  
    27  func (p *PollTask) Register(pb Pollable) {
    28  	p.m.Set(pb.ID(), pb)
    29  	pb.OnDone(func(_ bool, _ Task) {
    30  		p.Remove(pb)
    31  	})
    32  }
    33  
    34  func (p *PollTask) Remove(pb Pollable) {
    35  	p.m.Del(pb.ID())
    36  }
    37  
    38  func (p *PollTask) PollOne(id string) (bool, error) {
    39  	pb, ok := p.m.Get(id)
    40  	if !ok {
    41  		return false, ErrPollableEventsNotExists
    42  	}
    43  	return pb.IsDone(), nil
    44  }
    45  
    46  func (p *PollTask) Poll(f func(string, Pollable)) {
    47  	p.m.ForEach(func(id string, pb Pollable) bool {
    48  		if pb.IsDone() {
    49  			f(id, pb)
    50  		}
    51  		return true
    52  	})
    53  }
    54  
    55  func (p *PollTask) Callback(id string, f Finalizer) error {
    56  	pb, ok := p.m.Get(id)
    57  	if !ok {
    58  		return ErrPollableEventsNotExists
    59  	}
    60  	pb.OnDone(f)
    61  	return nil
    62  }
    63  
    64  func (p *PollTask) Kill(id string) error {
    65  	pb, ok := p.m.Get(id)
    66  	if !ok {
    67  		return ErrPollableEventsNotExists
    68  	}
    69  	pb.Interrupt()
    70  	return nil
    71  }
    72  
    73  func (p *PollTask) Get(id string) (Pollable, bool) {
    74  	return p.m.Get(id)
    75  }
    76  
    77  func (p *PollTask) ForEach(f func(string, Pollable) bool) {
    78  	p.m.ForEach(f)
    79  }