gitee.com/sy_183/go-common@v1.0.5-0.20231205030221-958cfe129b47/lifecycle.v2/task/task-executor.go (about)

     1  package task
     2  
     3  import (
     4  	"errors"
     5  	"gitee.com/sy_183/go-common/lifecycle.v2"
     6  	"sync/atomic"
     7  )
     8  
     9  var ErrTaskExecutorClosed = errors.New("task executor closed")
    10  
    11  type TaskExecutor struct {
    12  	lifecycle.Lifecycle
    13  	tasks       chan Task
    14  	interrupted atomic.Bool
    15  }
    16  
    17  func NewTaskExecutor(maxTask int) *TaskExecutor {
    18  	e := &TaskExecutor{
    19  		tasks: make(chan Task, maxTask),
    20  	}
    21  	e.Lifecycle = lifecycle.NewInterruptedWithFunc(nil, e.run)
    22  	return e
    23  }
    24  
    25  func (e *TaskExecutor) Init(maxTask int) {
    26  	e.tasks = make(chan Task, maxTask)
    27  	e.Lifecycle = lifecycle.NewInterruptedWithFunc(nil, e.run)
    28  }
    29  
    30  func (e *TaskExecutor) run(_ lifecycle.Lifecycle, interrupter chan struct{}) error {
    31  	for {
    32  		select {
    33  		case task := <-e.tasks:
    34  			if task.Do(interrupter) {
    35  				close(e.tasks)
    36  				return nil
    37  			}
    38  		case <-interrupter:
    39  			close(e.tasks)
    40  			interrupter <- struct{}{}
    41  			for task := range e.tasks {
    42  				if task.Do(interrupter) {
    43  					return nil
    44  				}
    45  			}
    46  			return nil
    47  		}
    48  	}
    49  }
    50  
    51  func (e *TaskExecutor) RunFunc() lifecycle.InterruptedRunFunc {
    52  	return e.run
    53  }
    54  
    55  func noPanicError() {
    56  	if e := recover(); e != nil {
    57  		if _, is := e.(error); !is {
    58  			panic(e)
    59  		}
    60  	}
    61  }
    62  
    63  func (e *TaskExecutor) noPanicPush(ch chan Task, task Task) (ok bool) {
    64  	defer noPanicError()
    65  	ch <- task
    66  	return true
    67  }
    68  
    69  func (e *TaskExecutor) noPanicTryPush(ch chan Task, task Task) (ok bool) {
    70  	defer noPanicError()
    71  	select {
    72  	case ch <- task:
    73  		return true
    74  	default:
    75  		return false
    76  	}
    77  }
    78  
    79  func (e *TaskExecutor) Do(task Task) (ok bool) {
    80  	return e.noPanicPush(e.tasks, task)
    81  }
    82  
    83  func (e *TaskExecutor) Try(task Task) (ok bool) {
    84  	return e.noPanicTryPush(e.tasks, task)
    85  }
    86  
    87  func (e *TaskExecutor) Len() int {
    88  	return len(e.tasks)
    89  }
    90  
    91  func (e *TaskExecutor) Cap() int {
    92  	return cap(e.tasks)
    93  }