go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/sync/promise/promise.go (about)

     1  // Copyright 2015 The LUCI Authors.
     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  	"errors"
    20  	"sync"
    21  )
    22  
    23  var (
    24  	// ErrNoData is an error returned when a request completes without available data.
    25  	ErrNoData = errors.New("promise: No Data")
    26  )
    27  
    28  // Generator is the Promise's generator type.
    29  type Generator func(context.Context) (any, error)
    30  
    31  // Promise is a promise structure with goroutine-safe methods that is
    32  // responsible for owning a single piece of data. Promises have multiple readers
    33  // and a single writer.
    34  //
    35  // Readers will retrieve the Promise's data via Get(). If the data has not been
    36  // populated, the reader will block pending the data. Once the data has been
    37  // delivered, all readers will unblock and receive a reference to the Promise's
    38  // data.
    39  type Promise struct {
    40  	signalC chan struct{} // Channel whose closing signals that the data is available.
    41  
    42  	// onGet, if not nil, is invoked when Get is called.
    43  	onGet func(context.Context)
    44  
    45  	data any // The Promise's data.
    46  	err  error       // The error status.
    47  }
    48  
    49  // New instantiates a new, empty Promise instance. The Promise's value will be
    50  // the value returned by the supplied generator function.
    51  //
    52  // The generator will be invoked immediately in its own goroutine.
    53  func New(ctx context.Context, gen Generator) *Promise {
    54  	p := Promise{
    55  		signalC: make(chan struct{}),
    56  	}
    57  
    58  	// Execute our generator function in a separate goroutine.
    59  	go p.runGen(ctx, gen)
    60  	return &p
    61  }
    62  
    63  // NewDeferred instantiates a new, empty Promise instance. The Promise's value
    64  // will be the value returned by the supplied generator function.
    65  //
    66  // Unlike New, the generator function will not be immediately executed. Instead,
    67  // it will be run when the first call to Get is made, and will use one of the
    68  // Get callers' goroutines.
    69  // goroutine a Get caller.
    70  func NewDeferred(gen Generator) *Promise {
    71  	var startOnce sync.Once
    72  
    73  	p := Promise{
    74  		signalC: make(chan struct{}),
    75  	}
    76  	p.onGet = func(ctx context.Context) {
    77  		startOnce.Do(func() { p.runGen(ctx, gen) })
    78  	}
    79  	return &p
    80  }
    81  
    82  func (p *Promise) runGen(ctx context.Context, gen Generator) {
    83  	defer close(p.signalC)
    84  	p.data, p.err = gen(ctx)
    85  }
    86  
    87  // Get returns the promise's value. If the value isn't set, Get will block until
    88  // the value is available, following the Context's timeout parameters.
    89  //
    90  // If the value is available, it will be returned with its error status. If the
    91  // context times out or is cancelled, the appropriate context error will be
    92  // returned.
    93  func (p *Promise) Get(ctx context.Context) (any, error) {
    94  	// If we have an onGet function, run it (deferred case).
    95  	if p.onGet != nil {
    96  		p.onGet(ctx)
    97  	}
    98  
    99  	// Block until at least one of these conditions is satisfied. If both are,
   100  	// "select" will choose one pseudo-randomly.
   101  	select {
   102  	case <-p.signalC:
   103  		return p.data, p.err
   104  
   105  	case <-ctx.Done():
   106  		// Make sure we don't actually have data.
   107  		select {
   108  		case <-p.signalC:
   109  			return p.data, p.err
   110  
   111  		default:
   112  			return nil, ctx.Err()
   113  		}
   114  	}
   115  }
   116  
   117  // Peek returns the promise's current value. If the value isn't set, Peek will
   118  // return immediately with ErrNoData.
   119  func (p *Promise) Peek() (any, error) {
   120  	select {
   121  	case <-p.signalC:
   122  		return p.data, p.err
   123  
   124  	default:
   125  		return nil, ErrNoData
   126  	}
   127  }