github.com/dtroyer-salad/og2/v2@v2.0.0-20240412154159-c47231610877/internal/syncutil/once.go (about)

     1  /*
     2  Copyright The ORAS Authors.
     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  
    16  package syncutil
    17  
    18  import (
    19  	"context"
    20  	"sync"
    21  	"sync/atomic"
    22  )
    23  
    24  // Once is an object that will perform exactly one action.
    25  // Unlike sync.Once, this Once allows the action to have return values.
    26  type Once struct {
    27  	result interface{}
    28  	err    error
    29  	status chan bool
    30  }
    31  
    32  // NewOnce creates a new Once instance.
    33  func NewOnce() *Once {
    34  	status := make(chan bool, 1)
    35  	status <- true
    36  	return &Once{
    37  		status: status,
    38  	}
    39  }
    40  
    41  // Do calls the function f if and only if Do is being called first time or all
    42  // previous function calls are cancelled, deadline exceeded, or panicking.
    43  // When `once.Do(ctx, f)` is called multiple times, the return value of the
    44  // first call of the function f is stored, and is directly returned for other
    45  // calls.
    46  // Besides the return value of the function f, including the error, Do returns
    47  // true if the function f passed is called first and is not cancelled, deadline
    48  // exceeded, or panicking. Otherwise, returns false.
    49  func (o *Once) Do(ctx context.Context, f func() (interface{}, error)) (bool, interface{}, error) {
    50  	defer func() {
    51  		if r := recover(); r != nil {
    52  			o.status <- true
    53  			panic(r)
    54  		}
    55  	}()
    56  	for {
    57  		select {
    58  		case inProgress := <-o.status:
    59  			if !inProgress {
    60  				return false, o.result, o.err
    61  			}
    62  			result, err := f()
    63  			if err == context.Canceled || err == context.DeadlineExceeded {
    64  				o.status <- true
    65  				return false, nil, err
    66  			}
    67  			o.result, o.err = result, err
    68  			close(o.status)
    69  			return true, result, err
    70  		case <-ctx.Done():
    71  			return false, nil, ctx.Err()
    72  		}
    73  	}
    74  }
    75  
    76  // OnceOrRetry is an object that will perform exactly one success action.
    77  type OnceOrRetry struct {
    78  	done atomic.Bool
    79  	lock sync.Mutex
    80  }
    81  
    82  // OnceOrRetry calls the function f if and only if Do is being called for the
    83  // first time for this instance of Once or all previous calls to Do are failed.
    84  func (o *OnceOrRetry) Do(f func() error) error {
    85  	// fast path
    86  	if o.done.Load() {
    87  		return nil
    88  	}
    89  
    90  	// slow path
    91  	o.lock.Lock()
    92  	defer o.lock.Unlock()
    93  
    94  	if o.done.Load() {
    95  		return nil
    96  	}
    97  	if err := f(); err != nil {
    98  		return err
    99  	}
   100  	o.done.Store(true)
   101  	return nil
   102  }