github.com/aaabigfish/gopkg@v1.1.0/mq/asynq/reader.go (about)

     1  package asynq
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/hibiken/asynq"
     7  )
     8  
     9  type Reader interface {
    10  	GetReader() *asynq.Server
    11  	AddHook(string, HandlerFunc)
    12  	AddHooks(Handler)
    13  	Run() error
    14  	Close()
    15  }
    16  
    17  type reader struct {
    18  	c      *RedisClientOpt
    19  	reader *asynq.Server
    20  	mux    *asynq.ServeMux
    21  	hooks  Handler
    22  	rwl    sync.RWMutex
    23  }
    24  
    25  func NewReader(opt *RedisClientOpt, c ...asynq.Config) Reader {
    26  	cfg := asynq.Config{}
    27  
    28  	if len(c) > 0 {
    29  		cfg = c[0]
    30  	}
    31  	r := asynq.NewServer(opt, cfg)
    32  
    33  	return &reader{
    34  		c:      opt,
    35  		reader: r,
    36  		mux:    asynq.NewServeMux(),
    37  		hooks:  make(map[string]HandlerFunc, 10),
    38  	}
    39  }
    40  
    41  func (r *reader) GetReader() *asynq.Server {
    42  	return r.reader
    43  }
    44  
    45  func (r *reader) AddHook(pattern string, hookFunc HandlerFunc) {
    46  	r.rwl.Lock()
    47  	r.hooks[pattern] = hookFunc
    48  	r.rwl.Unlock()
    49  }
    50  
    51  func (r *reader) AddHooks(hooks Handler) {
    52  	for pattern, hookFunc := range hooks {
    53  		if hookFunc == nil {
    54  			continue
    55  		}
    56  
    57  		r.rwl.Lock()
    58  		r.hooks[pattern] = hookFunc
    59  		r.rwl.Unlock()
    60  	}
    61  }
    62  
    63  func (r *reader) Run() error {
    64  	for pattern, handler := range r.hooks {
    65  		r.mux.HandleFunc(pattern, handler)
    66  	}
    67  
    68  	if err := r.reader.Run(r.mux); err != nil {
    69  		return err
    70  	}
    71  
    72  	return nil
    73  }
    74  
    75  func (r *reader) Close() {
    76  	if r.reader == nil {
    77  		r.reader.Shutdown()
    78  		return
    79  	}
    80  }