github.com/songzhibin97/gkit@v1.2.13/distributed/broker/broker.go (about)

     1  package broker
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/songzhibin97/gkit/distributed/retry"
     8  	"github.com/songzhibin97/gkit/options"
     9  )
    10  
    11  // registeredTask 注册的任务
    12  type registeredTask struct {
    13  	sync.RWMutex
    14  	item map[string]bool
    15  }
    16  
    17  // Register 注册
    18  func (r *registeredTask) Register(taskName string) {
    19  	r.Lock()
    20  	defer r.Unlock()
    21  	r.item[taskName] = true
    22  }
    23  
    24  // RegisterList 注册
    25  func (r *registeredTask) RegisterList(taskNameList ...string) {
    26  	r.Lock()
    27  	defer r.Unlock()
    28  	for _, taskName := range taskNameList {
    29  		r.item[taskName] = true
    30  	}
    31  }
    32  
    33  // Quit 注销
    34  func (r *registeredTask) Quit(taskName string) {
    35  	r.Lock()
    36  	defer r.Unlock()
    37  	delete(r.item, taskName)
    38  }
    39  
    40  // IsRegister 是否注册
    41  func (r *registeredTask) IsRegister(taskName string) bool {
    42  	r.RLock()
    43  	defer r.RUnlock()
    44  	return r.item[taskName]
    45  }
    46  
    47  // NewRegisteredTask 初始化任务注册器
    48  func NewRegisteredTask() *registeredTask {
    49  	return &registeredTask{
    50  		item: make(map[string]bool),
    51  	}
    52  }
    53  
    54  // Broker Broker
    55  type Broker struct {
    56  	// registeredTask 注册器
    57  	*registeredTask
    58  	// retry 是否重试
    59  	retry bool
    60  	// retryFn 重试函数
    61  	retryFn     func(ctx context.Context)
    62  	retryCtx    context.Context
    63  	retryCancel context.CancelFunc
    64  
    65  	stopCtx    context.Context
    66  	stopCancel context.CancelFunc
    67  }
    68  
    69  // NewBroker 初始化 Broker
    70  func NewBroker(r *registeredTask, ctx context.Context, options ...options.Option) *Broker {
    71  	b := &Broker{
    72  		registeredTask: r,
    73  	}
    74  	for _, option := range options {
    75  		option(b)
    76  	}
    77  	if b.retry == true && b.retryFn == nil {
    78  		b.retryFn = retry.Retry()
    79  	}
    80  	b.retryCtx, b.retryCancel = context.WithCancel(ctx)
    81  	b.stopCtx, b.stopCancel = context.WithCancel(ctx)
    82  	return b
    83  }
    84  
    85  func (b *Broker) GetRetry() bool {
    86  	return b.retry
    87  }
    88  
    89  func (b *Broker) GetRetryFn() func(ctx context.Context) {
    90  	return b.retryFn
    91  }
    92  
    93  func (b *Broker) GetRetryCtx() context.Context {
    94  	return b.retryCtx
    95  }
    96  
    97  func (b *Broker) GetStopCtx() context.Context {
    98  	return b.stopCtx
    99  }
   100  
   101  func (b *Broker) StopConsuming() {
   102  	b.retry = false
   103  	b.retryCancel()
   104  	b.stopCancel()
   105  }