github.com/pulumi/pulumi/sdk/v3@v3.108.1/go/common/promise/promise.go (about)

     1  // Copyright 2016-2023, Pulumi Corporation.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package promise
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"sync/atomic"
    21  
    22  	"github.com/pulumi/pulumi/sdk/v3/go/common/util/contract"
    23  )
    24  
    25  const (
    26  	statusUninitialized int32 = iota
    27  	statusPending
    28  	statusFulfilled
    29  	statusRejected
    30  )
    31  
    32  // Promise is a promise that can be resolved with a value of type T or rejected with an error. It is safe to call Result
    33  // on it multiple times from multiple goroutines. This is much more permissive than channels.
    34  type Promise[T any] struct {
    35  	done   chan struct{}
    36  	mutex  sync.Mutex
    37  	status atomic.Int32
    38  	result T
    39  	err    error
    40  }
    41  
    42  // Result waits for the promise to be resolved and returns the result.
    43  func (p *Promise[T]) Result(ctx context.Context) (T, error) {
    44  	if p.status.Load() == statusUninitialized {
    45  		panic("Promise must be initialized")
    46  	}
    47  
    48  	// Wait for either the promise or context to be done, if the context is done just exit with it's error
    49  	select {
    50  	case <-p.done:
    51  		break
    52  	case <-ctx.Done():
    53  		var t T
    54  		return t, ctx.Err()
    55  	}
    56  
    57  	contract.Assertf(p.status.Load() != statusPending, "Promise must be resolved")
    58  	// Only one of result or err will be set, the other will be the zero value so we can just return both.
    59  	return p.result, p.err
    60  }
    61  
    62  // TryResult returns the result and true if the promise has been resolved, otherwise it returns false.
    63  //
    64  //nolint:revive // This error _isn't_ an error from the function, so ignore the "error should be last" rule.
    65  func (p *Promise[T]) TryResult() (T, error, bool) {
    66  	// We don't need to lock here because we're just reading the status and the result and err are immutable
    67  	// once set.
    68  	status := p.status.Load()
    69  
    70  	if status == statusUninitialized {
    71  		panic("Promise must be initialized")
    72  	}
    73  
    74  	if status == statusPending {
    75  		var t T
    76  		return t, nil, false
    77  	}
    78  	// If the status is not pending then the promise is resolved and we can return the result and err. There
    79  	// is no race between status being set to fulfilled or rejected and result and err being changed.
    80  	return p.result, p.err, true
    81  }
    82  
    83  // CompletionSource is a source for a promise that can be resolved or rejected. It is safe to call Resolve or
    84  // Reject multiple times concurrently, the first will apply and all others will return that they couldn't set the
    85  // promise.
    86  type CompletionSource[T any] struct {
    87  	init    sync.Once
    88  	promise *Promise[T]
    89  }
    90  
    91  func (ps *CompletionSource[T]) Promise() *Promise[T] {
    92  	ps.init.Do(func() {
    93  		p := &Promise[T]{}
    94  		p.status.Store(statusPending)
    95  		p.done = make(chan struct{})
    96  		ps.promise = p
    97  	})
    98  	return ps.promise
    99  }
   100  
   101  func (ps *CompletionSource[T]) Fulfill(value T) bool {
   102  	promise := ps.Promise()
   103  	promise.mutex.Lock()
   104  	defer promise.mutex.Unlock()
   105  
   106  	contract.Assertf(promise.status.Load() != statusUninitialized, "Promise must be initialized")
   107  	if promise.status.Load() != statusPending {
   108  		return false
   109  	}
   110  	promise.result = value
   111  	promise.status.Store(statusFulfilled)
   112  	close(promise.done)
   113  	return true
   114  }
   115  
   116  func (ps *CompletionSource[T]) MustFulfill(value T) {
   117  	if !ps.Fulfill(value) {
   118  		panic("CompletionSource already resolved")
   119  	}
   120  }
   121  
   122  func (ps *CompletionSource[T]) Reject(err error) bool {
   123  	contract.Requiref(err != nil, "err", "err must not be nil")
   124  
   125  	promise := ps.Promise()
   126  	promise.mutex.Lock()
   127  	defer promise.mutex.Unlock()
   128  
   129  	contract.Assertf(promise.status.Load() != statusUninitialized, "Promise must be initialized")
   130  	if promise.status.Load() != statusPending {
   131  		return false
   132  	}
   133  	promise.err = err
   134  	promise.status.Store(statusRejected)
   135  	close(promise.done)
   136  	return true
   137  }
   138  
   139  func (ps *CompletionSource[T]) MustReject(err error) {
   140  	if !ps.Reject(err) {
   141  		panic("CompletionSource already resolved")
   142  	}
   143  }
   144  
   145  // Run runs the given function in a goroutine and returns a promise that will be resolved with the result of the
   146  // function.
   147  func Run[T any](f func() (T, error)) *Promise[T] {
   148  	ps := &CompletionSource[T]{}
   149  	go func() {
   150  		value, err := f()
   151  		if err != nil {
   152  			ps.Reject(err)
   153  		} else {
   154  			ps.Fulfill(value)
   155  		}
   156  	}()
   157  	return ps.Promise()
   158  }