github.com/lyft/flytestdlib@v0.3.12-0.20210213045714-8cdd111ecda1/futures/future.go (about)

     1  // This module implements a simple Async Futures for golang
     2  // Usage:
     3  // f := NewAsyncFuture(childCtx, func(ctx2 context.Context) (interface{}, error) {
     4  // can do large async / non blocking work
     5  //			return ...
     6  //		}
     7  // f.Ready() // can be checked for completion
     8  // f.Get() .. will block till the given sub-routine returns
     9  package futures
    10  
    11  import (
    12  	"context"
    13  	"fmt"
    14  	"sync"
    15  )
    16  
    17  // Provides a Future API for asynchronous completion of tasks
    18  type Future interface {
    19  	// Returns true if the Future is ready and either a value or error is available. Once Ready returns True, Get should return immediately
    20  	Ready() bool
    21  	// Get is a potentially blocking call, that returns the asynchronously computed value or an error
    22  	// If Get is called before Ready() returns True, then it will block till the future has been completed
    23  	Get(ctx context.Context) (interface{}, error)
    24  }
    25  
    26  // This is a synchronous future, where the values are available immediately on construction. This is used to maintain a synonymous API with both
    27  // Async and Sync tasks
    28  type SyncFuture struct {
    29  	// The actual value
    30  	val interface{}
    31  	// OR an error
    32  	err error
    33  }
    34  
    35  // Always returns true
    36  func (s SyncFuture) Ready() bool {
    37  	return true
    38  }
    39  
    40  // returns the previously provided value / error
    41  func (s *SyncFuture) Get(_ context.Context) (interface{}, error) {
    42  	return s.val, s.err
    43  }
    44  
    45  // Creates a new "completed" future that matches the async computation api
    46  func NewSyncFuture(val interface{}, err error) *SyncFuture {
    47  	return &SyncFuture{
    48  		val: val,
    49  		err: err,
    50  	}
    51  }
    52  
    53  // ErrAsyncFutureCanceled is returned when the async future is cancelled by invoking the cancel function on the context
    54  var ErrAsyncFutureCanceled = fmt.Errorf("async future was canceled")
    55  
    56  // An asynchronously completing future
    57  type AsyncFuture struct {
    58  	sync.Mutex
    59  	doneChannel chan bool
    60  	cancelFn    context.CancelFunc
    61  	// The actual value
    62  	val interface{}
    63  	// Or an error
    64  	err   error
    65  	ready bool
    66  }
    67  
    68  func (f *AsyncFuture) set(val interface{}, err error) {
    69  	f.Lock()
    70  	defer f.Unlock()
    71  	f.val = val
    72  	f.err = err
    73  	f.ready = true
    74  	f.doneChannel <- true
    75  }
    76  
    77  func (f *AsyncFuture) get() (interface{}, error) {
    78  	f.Lock()
    79  	defer f.Unlock()
    80  	return f.val, f.err
    81  }
    82  
    83  // returns whether the future is completed
    84  func (f *AsyncFuture) Ready() bool {
    85  	f.Lock()
    86  	defer f.Unlock()
    87  	return f.ready
    88  }
    89  
    90  // Returns results (interface{} or an error) OR blocks till the results are available.
    91  // If context is cancelled while waiting for results, an ErrAsyncFutureCanceled is returned
    92  func (f *AsyncFuture) Get(ctx context.Context) (interface{}, error) {
    93  	select {
    94  	case <-ctx.Done():
    95  		f.cancelFn()
    96  		return nil, ErrAsyncFutureCanceled
    97  	case <-f.doneChannel:
    98  		return f.get()
    99  	}
   100  }
   101  
   102  // Creates a new Async future, that will call the method `closure` and return the results from the execution of
   103  // this method
   104  func NewAsyncFuture(ctx context.Context, closure func(context.Context) (interface{}, error)) *AsyncFuture {
   105  	childCtx, cancel := context.WithCancel(ctx)
   106  	f := &AsyncFuture{
   107  		doneChannel: make(chan bool, 1),
   108  		cancelFn:    cancel,
   109  	}
   110  
   111  	go func(ctx2 context.Context, fut *AsyncFuture) {
   112  		val, err := closure(ctx2)
   113  		fut.set(val, err)
   114  	}(childCtx, f)
   115  	return f
   116  }