github.com/KarpelesLab/weak@v0.1.1/ref.go (about)

     1  package weak
     2  
     3  import (
     4  	"runtime"
     5  	"sync/atomic"
     6  	"unsafe"
     7  )
     8  
     9  // Ref is a weak reference to a Go object
    10  type Ref[T any] struct {
    11  	hidden uintptr
    12  	state  refState
    13  }
    14  
    15  func (wr *Ref[T]) value() *T {
    16  	v := atomic.LoadUintptr(&wr.hidden)
    17  	if v == 0 {
    18  		return nil
    19  	}
    20  	return (*T)(unsafe.Pointer(v))
    21  }
    22  
    23  // Get returns the value for a given weak reference pointer
    24  func (wr *Ref[T]) Get() *T {
    25  	for {
    26  		if wr.state.CaS(refALIVE, refINUSE) {
    27  			val := wr.value()
    28  			// all good
    29  			wr.state.Set(refALIVE) // set back to alive
    30  			return val
    31  		}
    32  		if wr.state.Get() == refDEAD {
    33  			return nil
    34  		}
    35  		runtime.Gosched()
    36  	}
    37  }
    38  
    39  // NewRef returns a reference to the object v that may be cleared by the garbage collector
    40  func NewRef[T any](v *T) *Ref[T] {
    41  	if v == nil {
    42  		return &Ref[T]{0, refDEAD}
    43  	}
    44  	wr := &Ref[T]{uintptr(unsafe.Pointer(v)), refALIVE}
    45  	var f func(p *T)
    46  	f = func(p *T) {
    47  		if wr.state.CaS(refALIVE, refDEAD) {
    48  			// we're now refdead, clear the pointer value
    49  			atomic.StoreUintptr(&wr.hidden, 0)
    50  			return
    51  		}
    52  		// this was not ALIVE, it means it was likely INUSE, re-set finalizer and wait
    53  		runtime.SetFinalizer(p, f)
    54  	}
    55  	runtime.SetFinalizer(v, f)
    56  
    57  	return wr
    58  }
    59  
    60  // NewRefDestroyer returns a reference to the object v that may be cleared by the garbage collector,
    61  // in which case destroy will be called.
    62  func NewRefDestroyer[T any](v *T, destroy func(v *T, wr *Ref[T])) *Ref[T] {
    63  	if v == nil {
    64  		return &Ref[T]{0, refDEAD}
    65  	}
    66  	wr := &Ref[T]{uintptr(unsafe.Pointer(v)), refALIVE}
    67  	var f func(p *T)
    68  	f = func(p *T) {
    69  		if wr.state.CaS(refALIVE, refDEAD) {
    70  			atomic.StoreUintptr(&wr.hidden, 0)
    71  			if destroy != nil {
    72  				go destroy(p, wr)
    73  			}
    74  			return
    75  		}
    76  		// this was not ALIVE, it means it was likely INUSE, re-set finalizer and wait
    77  		runtime.SetFinalizer(p, f)
    78  	}
    79  	runtime.SetFinalizer(v, f)
    80  
    81  	return wr
    82  }