github.com/weedge/lib@v0.0.0-20230424045628-a36dcc1d90e4/pool/asynctask/asynctask.go (about)

     1  package asynctask
     2  
     3  import (
     4  	"fmt"
     5  	"runtime/debug"
     6  	"sync"
     7  
     8  	"github.com/weedge/lib/container/set"
     9  	"github.com/weedge/lib/log"
    10  	"github.com/weedge/lib/strings"
    11  )
    12  
    13  var asyncTaskNames *set.HashSet
    14  
    15  func init() {
    16  	asyncTaskNames = set.NewSet()
    17  }
    18  
    19  type IAsyncTask interface {
    20  	Run() error
    21  }
    22  
    23  type AsyncTask struct {
    24  	name string
    25  	ch   chan IAsyncTask
    26  	wg   *sync.WaitGroup
    27  }
    28  
    29  func (this *AsyncTask) Close() {
    30  	asyncTaskNames.Remove(this.name)
    31  	close(this.ch)
    32  	this.wg.Wait()
    33  }
    34  
    35  func NewAsyncTask(name string, taskChanNumber int64, goNumber int, onError func(err error)) (*AsyncTask, error) {
    36  	if ok := asyncTaskNames.Contains(name); ok {
    37  		return nil, fmt.Errorf("asynctask name duplicated: %v", name)
    38  	}
    39  	asyncTaskNames.Add(name)
    40  
    41  	asyncTask := new(AsyncTask)
    42  	asyncTask.name = name
    43  	asyncTask.ch = make(chan IAsyncTask, taskChanNumber)
    44  	asyncTask.wg = &sync.WaitGroup{}
    45  	for i := 0; i < goNumber; i++ {
    46  		asyncTask.wg.Add(1)
    47  		asyncTask.run(onError)
    48  	}
    49  	return asyncTask, nil
    50  }
    51  
    52  func Recover() {
    53  	if e := recover(); e != nil {
    54  		log.Errorf("panic: %v, stack: %v", e, strings.BytesToString(debug.Stack()))
    55  	}
    56  }
    57  
    58  func (this *AsyncTask) run(onError func(err error)) {
    59  	go func(name string, ch chan IAsyncTask) {
    60  		defer this.wg.Done()
    61  		defer Recover()
    62  
    63  		log.Infof("AsyncTask [%v] created", name)
    64  		for {
    65  			asyncTask, ok := <-ch
    66  			if !ok {
    67  				break
    68  			}
    69  			if err := asyncTask.Run(); nil != err {
    70  				if nil != onError {
    71  					onError(err)
    72  				}
    73  			}
    74  		}
    75  		log.Infof("AsyncTask [%v] quit", name)
    76  	}(this.name, this.ch)
    77  }
    78  
    79  func (this *AsyncTask) Post(tasks ...IAsyncTask) {
    80  	for _, task := range tasks {
    81  		this.ch <- task
    82  	}
    83  }