go sync/weakmap

weakmap.go

package weakmap

import (
	"reflect"
	"runtime"
	"sync"
)

type Key = interface {}
type Value = interface {}

type WeakMap struct {
	m map[uintptr]Value
	l sync.RWMutex
}

func New() *WeakMap {
	return &WeakMap{
		m: make(map[uintptr]Value),
	}
}

// Key is weakly referenced.
func (w *WeakMap) Put(k Key, v Value) {
	runtime.SetFinalizer(k, w.finalizer)
	ptr := reflect.ValueOf(k).Pointer()
	w.l.Lock()
	w.m[ptr] = v
	w.l.Unlock()
}

func (w *WeakMap) Get(k Key) (Value, bool) {
	ptr := reflect.ValueOf(k).Pointer()
	w.l.RLock()
	v, ok := w.m[ptr]
	w.l.RUnlock()
	return v, ok
}

func (w *WeakMap) Len() int {
	w.l.RLock()
	v := len(w.m)
	w.l.RUnlock()
	return v
}

func (w *WeakMap) finalizer(k Key) {
	ptr := reflect.ValueOf(k).Pointer() // uintptr(unsafe.Pointer(&k))
	w.l.Lock()
	delete(w.m, ptr)
	w.l.Unlock()
}

weakmap_test.go

package weakmap

import (
	"github.com/stretchr/testify/assert"
	"runtime"
	"testing"
)

func TestWeakmapKeyWithoutRef(t *testing.T) {
	type E struct{ v string }
	var wm = New()
	{
		k := &E{"key"}
		v := &E{"value"}
		wm.Put(k, v)
		actualV, ok := wm.Get(k)
		assert.True(t, ok)
		assert.Equal(t, v, actualV)
		assert.Equal(t, 1, wm.Len())
	}
	runtime.GC()
	assert.Equal(t, 0, wm.Len())
}

func TestWeakmapKeyWithRef(t *testing.T) {
	type E struct{ v string }
	var wm = New()
	k := &E{"key"}
	{
		v := &E{"value"}
		wm.Put(k, v)
		actualV, ok := wm.Get(k)
		assert.True(t, ok)
		assert.Equal(t, v, actualV)
		assert.Equal(t, 1, wm.Len())
	}
	runtime.GC()
	assert.Equal(t, 1, wm.Len())
	capture(k.v)
}

//go:noinline
func capture(v interface{}) {}