oncefunc.go

     1  // Copyright 2022 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package sync
     6  
     7  // OnceFunc returns a function that invokes f only once. The returned function
     8  // may be called concurrently.
     9  //
    10  // If f panics, the returned function will panic with the same value on every call.
    11  func OnceFunc(f func()) func() {
    12  	// Use a struct so that there's a single heap allocation.
    13  	d := struct {
    14  		f     func()
    15  		once  Once
    16  		valid bool
    17  		p     any
    18  	}{
    19  		f: f,
    20  	}
    21  	return func() {
    22  		d.once.Do(func() {
    23  			defer func() {
    24  				d.f = nil // Do not keep f alive after invoking it.
    25  				d.p = recover()
    26  				if !d.valid {
    27  					// Re-panic immediately so on the first
    28  					// call the user gets a complete stack
    29  					// trace into f.
    30  					panic(d.p)
    31  				}
    32  			}()
    33  			d.f()
    34  			d.valid = true // Set only if f does not panic.
    35  		})
    36  		if !d.valid {
    37  			panic(d.p)
    38  		}
    39  	}
    40  }
    41  
    42  // OnceValue returns a function that invokes f only once and returns the value
    43  // returned by f. The returned function may be called concurrently.
    44  //
    45  // If f panics, the returned function will panic with the same value on every call.
    46  func OnceValue[T any](f func() T) func() T {
    47  	// Use a struct so that there's a single heap allocation.
    48  	d := struct {
    49  		f      func() T
    50  		once   Once
    51  		valid  bool
    52  		p      any
    53  		result T
    54  	}{
    55  		f: f,
    56  	}
    57  	return func() T {
    58  		d.once.Do(func() {
    59  			defer func() {
    60  				d.f = nil
    61  				d.p = recover()
    62  				if !d.valid {
    63  					panic(d.p)
    64  				}
    65  			}()
    66  			d.result = d.f()
    67  			d.valid = true
    68  		})
    69  		if !d.valid {
    70  			panic(d.p)
    71  		}
    72  		return d.result
    73  	}
    74  }
    75  
    76  // OnceValues returns a function that invokes f only once and returns the values
    77  // returned by f. The returned function may be called concurrently.
    78  //
    79  // If f panics, the returned function will panic with the same value on every call.
    80  func OnceValues[T1, T2 any](f func() (T1, T2)) func() (T1, T2) {
    81  	// Use a struct so that there's a single heap allocation.
    82  	d := struct {
    83  		f     func() (T1, T2)
    84  		once  Once
    85  		valid bool
    86  		p     any
    87  		r1    T1
    88  		r2    T2
    89  	}{
    90  		f: f,
    91  	}
    92  	return func() (T1, T2) {
    93  		d.once.Do(func() {
    94  			defer func() {
    95  				d.f = nil
    96  				d.p = recover()
    97  				if !d.valid {
    98  					panic(d.p)
    99  				}
   100  			}()
   101  			d.r1, d.r2 = d.f()
   102  			d.valid = true
   103  		})
   104  		if !d.valid {
   105  			panic(d.p)
   106  		}
   107  		return d.r1, d.r2
   108  	}
   109  }
   110
View as plain text