github.com/lingyao2333/mo-zero@v1.4.1/core/fx/timeout.go (about)

     1  package fx
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"runtime/debug"
     7  	"strings"
     8  	"time"
     9  )
    10  
    11  var (
    12  	// ErrCanceled is the error returned when the context is canceled.
    13  	ErrCanceled = context.Canceled
    14  	// ErrTimeout is the error returned when the context's deadline passes.
    15  	ErrTimeout = context.DeadlineExceeded
    16  )
    17  
    18  // DoOption defines the method to customize a DoWithTimeout call.
    19  type DoOption func() context.Context
    20  
    21  // DoWithTimeout runs fn with timeout control.
    22  func DoWithTimeout(fn func() error, timeout time.Duration, opts ...DoOption) error {
    23  	parentCtx := context.Background()
    24  	for _, opt := range opts {
    25  		parentCtx = opt()
    26  	}
    27  	ctx, cancel := context.WithTimeout(parentCtx, timeout)
    28  	defer cancel()
    29  
    30  	// create channel with buffer size 1 to avoid goroutine leak
    31  	done := make(chan error, 1)
    32  	panicChan := make(chan interface{}, 1)
    33  	go func() {
    34  		defer func() {
    35  			if p := recover(); p != nil {
    36  				// attach call stack to avoid missing in different goroutine
    37  				panicChan <- fmt.Sprintf("%+v\n\n%s", p, strings.TrimSpace(string(debug.Stack())))
    38  			}
    39  		}()
    40  		done <- fn()
    41  	}()
    42  
    43  	select {
    44  	case p := <-panicChan:
    45  		panic(p)
    46  	case err := <-done:
    47  		return err
    48  	case <-ctx.Done():
    49  		return ctx.Err()
    50  	}
    51  }
    52  
    53  // WithContext customizes a DoWithTimeout call with given ctx.
    54  func WithContext(ctx context.Context) DoOption {
    55  	return func() context.Context {
    56  		return ctx
    57  	}
    58  }