github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/kit/mq/mem_mq/task.go (about)

     1  package mem_mq
     2  
     3  import (
     4  	"container/list"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/pkg/errors"
    10  
    11  	"github.com/machinefi/w3bstream/pkg/depends/kit/mq"
    12  )
    13  
    14  func New(limit int) *TaskManager {
    15  	if limit == 0 {
    16  		limit = 256
    17  	}
    18  	return &TaskManager{
    19  		l:   list.New(),
    20  		m:   map[string]*list.Element{},
    21  		lmt: limit,
    22  		sig: make(chan struct{}, limit),
    23  	}
    24  }
    25  
    26  type TaskManager struct {
    27  	m   map[string]*list.Element
    28  	l   *list.List
    29  	lmt int
    30  	sig chan struct{}
    31  
    32  	rwm sync.RWMutex
    33  }
    34  
    35  var _ mq.TaskManager = (*TaskManager)(nil)
    36  
    37  func (tm *TaskManager) Push(ch string, t mq.Task) error {
    38  	tm.rwm.Lock()
    39  	tm.m[key(ch, t.ID())] = tm.l.PushBack(t)
    40  	tm.rwm.Unlock()
    41  
    42  	select {
    43  	case tm.sig <- struct{}{}:
    44  		return nil
    45  	case <-time.After(time.Second):
    46  		return errors.Wrap(
    47  			mq.ErrPushTaskTimeout,
    48  			fmt.Sprintf("len: %d capacity: %d", len(tm.sig), tm.lmt),
    49  		)
    50  	}
    51  }
    52  
    53  func (tm *TaskManager) Pop(ch string) (mq.Task, error) {
    54  	<-tm.sig
    55  
    56  	tm.rwm.Lock()
    57  	defer tm.rwm.Unlock()
    58  
    59  	elem := tm.l.Front()
    60  	if elem == nil {
    61  		return nil, nil
    62  	}
    63  	tm.l.Remove(elem)
    64  
    65  	t, ok := elem.Value.(mq.Task)
    66  	if !ok {
    67  		return nil, nil
    68  	}
    69  
    70  	k := key(ch, t.ID())
    71  	if _, ok = tm.m[k]; !ok {
    72  		return nil, nil
    73  	}
    74  	return t, tm.remove(k)
    75  }
    76  
    77  func (tm *TaskManager) Remove(ch string, id string) error {
    78  	tm.rwm.Lock()
    79  	defer tm.rwm.Unlock()
    80  
    81  	return tm.remove(key(ch, id))
    82  }
    83  
    84  func (tm *TaskManager) Clear(_ string) error {
    85  	*tm = *New(tm.lmt)
    86  	return nil
    87  }
    88  
    89  func (tm *TaskManager) remove(key string) error {
    90  	if t := tm.m[key]; t != nil {
    91  		tm.l.Remove(t)
    92  		delete(tm.m, key)
    93  	}
    94  	return nil
    95  }
    96  
    97  func key(ch, id string) string { return ch + "::" + id }