gitee.com/h79/goutils@v1.22.10/common/scheduler/pool.go (about)

     1  package scheduler
     2  
     3  import (
     4  	"gitee.com/h79/goutils/common/logger"
     5  	"gitee.com/h79/goutils/common/system"
     6  	"sync/atomic"
     7  	"time"
     8  )
     9  
    10  type IdleHandler interface {
    11  	Idle(num int32)
    12  }
    13  
    14  // IdleFunc idle working number
    15  type IdleFunc func(num int32)
    16  
    17  func (fn IdleFunc) Idle(num int32) {
    18  	fn(num)
    19  }
    20  
    21  type Pool struct {
    22  	jobs        chan *Job
    23  	workers     chan *jobWorker
    24  	stop        chan bool
    25  	delay       *Delay
    26  	tm          *time.Ticker
    27  	numWorkers  int32 //总数
    28  	runWorkers  int32 //正常运行
    29  	idleWorking int32 //空闲worker数
    30  	idleHandler IdleHandler
    31  	running     system.RunningCheck
    32  }
    33  
    34  func NewPool(numWorkers, jobQueueLen int) *Pool {
    35  	if numWorkers <= 0 || jobQueueLen <= 0 {
    36  		panic("length is zero")
    37  	}
    38  	var num = int32(numWorkers)
    39  	pool := &Pool{
    40  		jobs:        make(chan *Job, jobQueueLen),
    41  		workers:     make(chan *jobWorker, numWorkers),
    42  		stop:        make(chan bool),
    43  		delay:       NewDelay(),
    44  		tm:          time.NewTicker(time.Millisecond * 10),
    45  		numWorkers:  num,
    46  		runWorkers:  num,
    47  		idleWorking: num,
    48  	}
    49  	pool.createWorkers(numWorkers)
    50  	pool.Run()
    51  	return pool
    52  }
    53  
    54  func (p *Pool) SetIdleHandler(fn IdleHandler) *Pool {
    55  	p.idleHandler = fn
    56  	return p
    57  }
    58  
    59  func (p *Pool) WithIdleFunc(fn func(num int32)) *Pool {
    60  	return p.SetIdleHandler(IdleFunc(fn))
    61  }
    62  
    63  func (p *Pool) AddJob(job *Job) {
    64  	if job == nil {
    65  		return
    66  	}
    67  	if system.IsQuit() {
    68  		logger.W("Task", "application is quited,not can add for jobId= %s", job.GetId())
    69  		return
    70  	}
    71  	ww := atomic.LoadInt32(&p.runWorkers)
    72  	if ww <= p.numWorkers/2 {
    73  		// 中间有原因，线程池中的线程，由于某些原因，退出了，需要再启动
    74  		p.createWorkers(int(p.numWorkers - ww))
    75  	}
    76  	p.Run()
    77  	p.jobs <- job
    78  }
    79  
    80  func (p *Pool) Run() {
    81  	p.running.GoRunning(p.dispatch)
    82  }
    83  
    84  func (p *Pool) HasWorking() bool {
    85  	return atomic.LoadInt32(&p.idleWorking) != p.numWorkers
    86  }
    87  
    88  func (p *Pool) Release() {
    89  	system.Stop(time.Second, p.stop)
    90  	p.idleHandle(p.numWorkers)
    91  }
    92  
    93  func (p *Pool) createWorkers(num int) {
    94  	for i := 0; i < num; i++ {
    95  		worker := newWorker(p)
    96  		worker.start()
    97  	}
    98  }
    99  
   100  func (p *Pool) idleHandle(idleWork int32) {
   101  	if p.idleHandler != nil {
   102  		p.idleHandler.Idle(idleWork)
   103  	}
   104  }
   105  
   106  func (p *Pool) workerRunning() {
   107  	atomic.AddInt32(&p.runWorkers, 1)
   108  }
   109  
   110  func (p *Pool) workerQuit() {
   111  	atomic.AddInt32(&p.runWorkers, -1)
   112  }
   113  
   114  func (p *Pool) workerIdle() {
   115  	p.idleHandle(atomic.AddInt32(&p.idleWorking, 1))
   116  }
   117  
   118  func (p *Pool) dispatch() {
   119  	defer p.tm.Stop()
   120  	defer p.idleHandle(p.numWorkers)
   121  	for {
   122  		select {
   123  		case <-p.stop:
   124  			p.stopWorks(1)
   125  			p.stop <- true
   126  			return
   127  
   128  		case <-system.Closed():
   129  			p.stopWorks(2)
   130  			return
   131  
   132  		case job := <-p.jobs:
   133  			if job.Delay > 0 {
   134  				p.delayJob(job)
   135  			} else {
   136  				p.executeJob(job)
   137  			}
   138  		case <-p.tm.C:
   139  			p.checkDelayJob()
   140  		}
   141  	}
   142  }
   143  
   144  func (p *Pool) stopWorks(s int) {
   145  	for i := 0; i < len(p.workers); i++ {
   146  		worker := <-p.workers
   147  		worker.stop <- s
   148  		<-worker.stop
   149  	}
   150  	atomic.StoreInt32(&p.idleWorking, p.numWorkers)
   151  	// delay
   152  	opts := With(s)
   153  	p.delay.Exec(opts...)
   154  
   155  	for i := 0; i < len(p.jobs); i++ {
   156  		job, ok := <-p.jobs
   157  		if !ok {
   158  			return
   159  		}
   160  		_, _ = job.Execute(opts...)
   161  	}
   162  }
   163  
   164  func (p *Pool) idleWork(w *jobWorker) {
   165  	if system.IsQuit() {
   166  		return
   167  	}
   168  	p.workers <- w
   169  }
   170  
   171  func (p *Pool) executeJob(job *Job) {
   172  	if system.IsQuit() {
   173  		_, _ = job.Execute(WithQuited())
   174  		return
   175  	}
   176  	worker := <-p.workers
   177  	p.idleHandle(atomic.AddInt32(&p.idleWorking, -1))
   178  	worker.job <- job
   179  }
   180  
   181  func (p *Pool) delayJob(job *Job) {
   182  	p.delay.Add(job)
   183  }
   184  
   185  func (p *Pool) checkDelayJob() {
   186  	p.delay.Check(p)
   187  }