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

     1  package weak
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  // Map is a thread safe map for objects to be kept as weak references, useful for cache/etc
     8  type Map[K comparable, T any] struct {
     9  	m map[K]*Ref[T]
    10  	l sync.RWMutex
    11  }
    12  
    13  // Object implementing Destroyable added to a Map will have their Destroy() method called
    14  // when the object is about to be removed. This replaces use of a finalizer.
    15  type Destroyable interface {
    16  	Destroy()
    17  }
    18  
    19  // NewMap returns a new weak reference map
    20  func NewMap[K comparable, T any]() *Map[K, T] {
    21  	res := &Map[K, T]{
    22  		m: make(map[K]*Ref[T]),
    23  	}
    24  
    25  	return res
    26  }
    27  
    28  // Get returns the value at index k in the map. If no such value exists, nil is returned.
    29  func (w *Map[K, T]) Get(k K) *T {
    30  	w.l.RLock()
    31  	defer w.l.RUnlock()
    32  
    33  	v, ok := w.m[k]
    34  	if !ok {
    35  		return nil
    36  	}
    37  
    38  	return v.Get()
    39  }
    40  
    41  // Set inserts the value v if it does not already exists, and return it. If a value v
    42  // already exists, then the previous value is returned.
    43  func (w *Map[K, T]) Set(k K, v *T) *T {
    44  	w.l.Lock()
    45  	defer w.l.Unlock()
    46  
    47  	// already exists?
    48  	wr, ok := w.m[k]
    49  	if ok {
    50  		v2 := wr.Get()
    51  		if v2 != nil {
    52  			// return past (still alive) value
    53  			return v2
    54  		}
    55  	}
    56  
    57  	// store new value
    58  	wr = NewRefDestroyer(v, func(dv *T, wr *Ref[T]) {
    59  		w.destroy(wr, dv, k)
    60  	})
    61  	w.m[k] = wr
    62  
    63  	return v
    64  }
    65  
    66  // Delete removes element at key k from the map. This doesn't call Destroy immediately
    67  // as this would typically happen when the object is actually cleared by the garbage
    68  // collector and instances of said object may still be used.
    69  func (w *Map[K, T]) Delete(k K) {
    70  	w.l.Lock()
    71  	defer w.l.Unlock()
    72  
    73  	delete(w.m, k)
    74  }
    75  
    76  func (w *Map[K, T]) destroy(wr *Ref[T], ptr *T, k K) {
    77  	w.l.Lock()
    78  	defer w.l.Unlock()
    79  
    80  	wr2, ok := w.m[k]
    81  	if !ok {
    82  		return
    83  	}
    84  	if wr == wr2 {
    85  		delete(w.m, k)
    86  	}
    87  
    88  	if v, ok := any(ptr).(Destroyable); ok {
    89  		go v.Destroy()
    90  	}
    91  }