Skip to content

Instantly share code, notes, and snippets.

@shyiko
Created June 16, 2019 04:05
Show Gist options
  • Save shyiko/e6e2a67e1c882a5034a3a290f9cc1a31 to your computer and use it in GitHub Desktop.
Save shyiko/e6e2a67e1c882a5034a3a290f9cc1a31 to your computer and use it in GitHub Desktop.
go sync/weakmap
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()
}
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{}) {}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment