github.com/songzhibin97/gkit@v1.2.13/distributed/task/task.go (about)

     1  package task
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"reflect"
     7  )
     8  
     9  type signatureCtxType struct{}
    10  
    11  var (
    12  	ErrDispatching = errors.New("dispatch task err")
    13  	signatureCtx   signatureCtxType
    14  )
    15  
    16  type Task struct {
    17  	// TaskFunc 执行任务的函数
    18  	TaskFunc reflect.Value
    19  	// UseContext 是否使用上下文
    20  	UseContext bool
    21  	// Context 上下文信息
    22  	Context context.Context
    23  	// Args 执行任务需要的参数
    24  	Args []reflect.Value
    25  }
    26  
    27  // TransformArgs 将[]Arg转化为[]reflect.Value
    28  func (t *Task) TransformArgs(args []Arg) error {
    29  	argValues := make([]reflect.Value, len(args))
    30  
    31  	for i, arg := range args {
    32  		argValue, err := ReflectValue(arg.Type, arg.Value)
    33  		if err != nil {
    34  			return err
    35  		}
    36  		argValues[i] = argValue
    37  	}
    38  
    39  	t.Args = argValues
    40  	return nil
    41  }
    42  
    43  // Call 调用方法
    44  func (t *Task) Call() (taskResults []*Result, err error) {
    45  	// 防止意外panic
    46  	defer func() {
    47  		if e := recover(); e != nil {
    48  			switch er := e.(type) {
    49  			default:
    50  				err = ErrDispatching
    51  			case error:
    52  				err = er
    53  			case string:
    54  				err = errors.New(er)
    55  			}
    56  		}
    57  	}()
    58  
    59  	args := t.Args
    60  	if t.UseContext {
    61  		ctxValue := reflect.ValueOf(t.Context)
    62  		args = append([]reflect.Value{ctxValue}, args...)
    63  	}
    64  
    65  	// 调用任务
    66  	results := t.TaskFunc.Call(args)
    67  
    68  	if len(results) == 0 {
    69  		return nil, ErrTaskReturnNoValue
    70  	}
    71  	// 按照规定最后一个参数是 err
    72  	lastResult := results[len(results)-1]
    73  	if !lastResult.IsNil() {
    74  		// err 不为nil
    75  
    76  		// 如果该错误实现了 Retrievable 接口
    77  		if lastResult.Type().Implements(retrievableInterface) {
    78  			return nil, lastResult.Interface().(ErrRetryTaskLater)
    79  		}
    80  		// 如果该错实现了 error 接口
    81  		if lastResult.Type().Implements(errInterface) {
    82  			return nil, lastResult.Interface().(error)
    83  		}
    84  		// 如果最后一个返回值没有满足error接口
    85  		return nil, ErrTaskReturnNoErr
    86  	}
    87  	taskResults = make([]*Result, 0, len(results)-1)
    88  	for i := 0; i < len(results)-1; i++ {
    89  		val := results[i].Interface()
    90  		typeStr := reflect.TypeOf(val).String()
    91  		taskResults = append(taskResults, &Result{
    92  			Type:  typeStr,
    93  			Value: val,
    94  		})
    95  	}
    96  	return taskResults, err
    97  }
    98  
    99  // SignatureFromContext 获取上下文任务签名
   100  func SignatureFromContext(ctx context.Context) *Signature {
   101  	if ctx == nil {
   102  		return nil
   103  	}
   104  	v := ctx.Value(signatureCtx)
   105  	if v == nil {
   106  		return nil
   107  	}
   108  	signature, _ := v.(*Signature)
   109  	return signature
   110  }
   111  
   112  // NewTaskWithSignature 初始化Task通过Signature
   113  func NewTaskWithSignature(taskFunc interface{}, signature *Signature) (*Task, error) {
   114  	ctx := context.WithValue(context.Background(), signatureCtx, signature)
   115  	task := &Task{
   116  		TaskFunc: reflect.ValueOf(taskFunc),
   117  		Context:  ctx,
   118  	}
   119  	taskFuncType := reflect.TypeOf(taskFunc)
   120  	if taskFuncType.NumIn() > 0 {
   121  		if taskFuncType.In(0) == ctxTypeInterface {
   122  			task.UseContext = true
   123  		}
   124  	}
   125  	if err := task.TransformArgs(signature.Args); err != nil {
   126  		return nil, err
   127  	}
   128  	return task, nil
   129  }
   130  
   131  // NewTask 初始化Task
   132  func NewTask(taskFunc interface{}, args []Arg) (*Task, error) {
   133  	task := &Task{
   134  		TaskFunc: reflect.ValueOf(taskFunc),
   135  		Context:  context.Background(),
   136  	}
   137  	taskFuncType := reflect.TypeOf(taskFunc)
   138  	if taskFuncType.NumIn() > 0 {
   139  		if taskFuncType.In(0) == ctxTypeInterface {
   140  			task.UseContext = true
   141  		}
   142  	}
   143  	if err := task.TransformArgs(args); err != nil {
   144  		return nil, err
   145  	}
   146  	return task, nil
   147  }