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

     1  package router
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/MeteorsLiu/rand"
     7  	"github.com/MeteorsLiu/simpleMQ/queue"
     8  	"github.com/MeteorsLiu/simpleMQ/worker"
     9  	"github.com/alphadose/haxmap"
    10  )
    11  
    12  const (
    13  	DefaultWorkerSize      = 100
    14  	DefaultWorkerSpwanSize = 10
    15  	DefaultWorkerCap       = 10
    16  )
    17  
    18  type Options func(*Router)
    19  
    20  func WithNewQueueFunc(f queue.NewQueue) Options {
    21  	return func(r *Router) {
    22  		r.newQueueFunc = f
    23  	}
    24  }
    25  
    26  func WithWorkerGroupCap(cap int) Options {
    27  	return func(r *Router) {
    28  		r.workersCap = cap
    29  	}
    30  }
    31  
    32  func WithWorkerSize(size int) Options {
    33  	return func(r *Router) {
    34  		r.workerSize = size
    35  	}
    36  }
    37  
    38  func WithWorkerSpwanSize(spwan int) Options {
    39  	return func(r *Router) {
    40  		r.workerSpawnSize = spwan
    41  	}
    42  }
    43  
    44  type Router struct {
    45  	sync.RWMutex
    46  	workerSize      int
    47  	workerSpawnSize int
    48  	workersCap      int
    49  	workers         []*worker.Worker
    50  	newQueueFunc    queue.NewQueue
    51  	routerPath      *haxmap.Map[string, *worker.Worker]
    52  }
    53  
    54  func NewRouter(opts ...Options) *Router {
    55  	r := &Router{
    56  		workerSize:      DefaultWorkerSize,
    57  		workerSpawnSize: DefaultWorkerSpwanSize,
    58  		newQueueFunc:    queue.NewSimpleQueue,
    59  		workersCap:      DefaultWorkerCap,
    60  		routerPath:      haxmap.New[string, *worker.Worker](),
    61  	}
    62  	for _, o := range opts {
    63  		o(r)
    64  	}
    65  	// new one for the task
    66  	r.setupWorker()
    67  	return r
    68  }
    69  
    70  func (r *Router) setupWorker() *worker.Worker {
    71  	r.Lock()
    72  	defer r.Unlock()
    73  	if len(r.workers) >= r.workersCap && r.workersCap > 0 {
    74  		return nil
    75  	}
    76  	newWorker := worker.NewWorker(r.workerSize, r.workerSpawnSize, r.newQueueFunc(), true)
    77  	r.workers = append(r.workers, newWorker)
    78  
    79  	return newWorker
    80  }
    81  
    82  func (r *Router) getWorkerBy(name string) *worker.Worker {
    83  	routerWorker, ok := r.routerPath.Get(name)
    84  	if !ok {
    85  		routerWorker = worker.NewWorker(r.workerSize, r.workerSpawnSize, r.newQueueFunc(), true)
    86  		r.routerPath.Set(name, routerWorker)
    87  	}
    88  	return routerWorker
    89  }
    90  
    91  func (r *Router) findWorker() *worker.Worker {
    92  	if tryNew := r.setupWorker(); tryNew != nil {
    93  		return tryNew
    94  	}
    95  	r.RLock()
    96  	defer r.RUnlock()
    97  	return r.workers[rand.Intn(len(r.workers))]
    98  }
    99  
   100  func (r *Router) killTask(id string) {
   101  	r.RLock()
   102  	defer r.RUnlock()
   103  
   104  	for _, worker := range r.workers {
   105  		if err := worker.KillTask(id); err == nil {
   106  			return
   107  		}
   108  	}
   109  }
   110  
   111  func (r *Router) Dispatch(f func() error, callback ...queue.Finalizer) queue.Task {
   112  	task := queue.NewTask(f)
   113  	r.findWorker().Publish(task, callback...)
   114  	return task
   115  }
   116  
   117  func (r *Router) DispatchTask(f queue.Task) {
   118  	r.findWorker().Publish(f)
   119  }
   120  
   121  func (r *Router) DispatchPath(path string, f func() error, callback ...queue.Finalizer) queue.Task {
   122  	task := queue.NewTask(f)
   123  	r.getWorkerBy(path).Publish(task, callback...)
   124  	return task
   125  }
   126  
   127  func (r *Router) DispatchPathTask(path string, task queue.Task) {
   128  	r.getWorkerBy(path).Publish(task)
   129  }
   130  
   131  func (r *Router) StopByID(id string) {
   132  	r.killTask(id)
   133  }
   134  
   135  func (r *Router) Stop() {
   136  	r.RLock()
   137  	for _, w := range r.workers {
   138  		w.Stop()
   139  	}
   140  	r.RUnlock()
   141  
   142  	r.routerPath.ForEach(func(_ string, w *worker.Worker) bool {
   143  		w.Stop()
   144  		return true
   145  	})
   146  }