github.com/codingeasygo/util@v0.0.0-20231206062002-1ce2f004b7d9/task/task.go (about)

     1  package task
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"github.com/codingeasygo/util/xdebug"
     8  )
     9  
    10  var Shared = NewRunner()
    11  
    12  func Call(state interface{}, call func(state interface{})) {
    13  	Shared.Call(state, call)
    14  }
    15  
    16  type Task struct {
    17  	State  interface{}
    18  	Caller func(state interface{})
    19  }
    20  
    21  type Runner struct {
    22  	max     int64
    23  	running int64
    24  	queue   chan *Task
    25  	waiter  sync.WaitGroup
    26  	locker  sync.RWMutex
    27  }
    28  
    29  func NewRunner() (runner *Runner) {
    30  	runner = &Runner{
    31  		queue:  make(chan *Task, 1024),
    32  		waiter: sync.WaitGroup{},
    33  		locker: sync.RWMutex{},
    34  		max:    3,
    35  	}
    36  	return
    37  }
    38  
    39  func (r *Runner) Max(max int) {
    40  	r.max = int64(max)
    41  }
    42  
    43  func (r *Runner) Stop() {
    44  	r.locker.Lock()
    45  	for i := int64(0); i < r.running; i++ {
    46  		r.queue <- nil
    47  	}
    48  	r.running = 0
    49  	r.locker.Unlock()
    50  	r.waiter.Wait()
    51  }
    52  
    53  func (r *Runner) run() {
    54  	defer r.waiter.Done()
    55  	running := true
    56  	for running {
    57  		t := <-r.queue
    58  		if t == nil {
    59  			running = false
    60  			break
    61  		}
    62  		func(t_ *Task) {
    63  			defer func() {
    64  				err := recover()
    65  				if err != nil {
    66  					fmt.Printf("Runner call panic with %v, callstack is \n%v", err, xdebug.CallStack())
    67  				}
    68  			}()
    69  			t_.Caller(t_.State)
    70  		}(t)
    71  		r.locker.Lock()
    72  		if r.running > r.max {
    73  			r.running--
    74  			running = false
    75  		}
    76  		r.locker.Unlock()
    77  	}
    78  }
    79  
    80  func (r *Runner) Call(state interface{}, call func(state interface{})) {
    81  	r.locker.Lock()
    82  	if r.running < r.max {
    83  		r.running++
    84  		r.waiter.Add(1)
    85  		go r.run()
    86  	}
    87  	r.locker.Unlock()
    88  	r.queue <- &Task{State: state, Caller: call}
    89  }