Changed WeakMap implementation to avoid memory leaks in some common usage scenarios. Fixes #250.

README.md

@@ -27,19 +27,31 @@ Known incompatibilities and caveats
 -----------------------------------
 
 ### WeakMap
-WeakMap maintains "hard" references to its values. This means if a value references a key in a WeakMap or a WeakMap
-itself, it will not be garbage-collected until the WeakMap becomes unreferenced. To illustrate this:
+WeakMap is implemented by embedding references to the values into the keys. This means that as long
+as the key is reachable all values associated with it in any weak maps also remain reachable and therefore
+cannot be garbage collected even if they are not otherwise referenced, even after the WeakMap is gone.
+The reference to the value is dropped either when the key is explicitly removed from the WeakMap or when the
+key becomes unreachable.
 
-```go
+To illustrate this:
+
+```javascript
 var m = new WeakMap();
 var key = {};
-m.set(key, {key: key});
-// or m.set(key, key);
-key = undefined; // The value will NOT become garbage-collectable at this point
-m = undefined; // But it will at this point.
+var value = {/* a very large object */};
+m.set(key, value);
+value = undefined;
+m = undefined; // The value does NOT become garbage-collectable at this point
+key = undefined; // Now it does
+// m.delete(key); // This would work too
 ```
 
-Note, this does not have any effect on the application logic, but causes a higher-than-expected memory usage.
+The reason for it is the limitation of the Go runtime. At the time of writing (version 1.15) having a finalizer
+set on an object which is part of a reference cycle makes the whole cycle non-garbage-collectable. The solution
+above is the only reasonable way I can think of without involving finalizers. This is the third attempt
+(see https://github.com/dop251/goja/issues/250 and https://github.com/dop251/goja/issues/199 for more details).
+
+Note, this does not have any effect on the application logic, but may cause a higher-than-expected memory usage.
 
 FAQ
 ---

builtin_weakmap.go

@@ -1,63 +1,35 @@
 package goja
 
-type weakMap struct {
-	data map[uint64]Value
-}
+type weakMap uint64
 
 type weakMapObject struct {
 	baseObject
-	m *weakMap
-}
-
-func newWeakMap() *weakMap {
-	return &weakMap{
-		data: make(map[uint64]Value),
-	}
+	m weakMap
 }
 
 func (wmo *weakMapObject) init() {
 	wmo.baseObject.init()
-	wmo.m = newWeakMap()
-}
-
-func (wm *weakMap) removeId(id uint64) {
-	delete(wm.data, id)
+	wmo.m = weakMap(wmo.val.runtime.genId())
 }
 
-func (wm *weakMap) set(key *Object, value Value) {
-	ref := key.getWeakRef()
-	wm.data[ref.id] = value
-	key.runtime.addWeakKey(ref.id, wm)
+func (wm weakMap) set(key *Object, value Value) {
+	key.getWeakRefs()[wm] = value
 }
 
-func (wm *weakMap) get(key *Object) Value {
-	ref := key.weakRef
-	if ref == nil {
-		return nil
-	}
-	ret := wm.data[ref.id]
-	return ret
+func (wm weakMap) get(key *Object) Value {
+	return key.weakRefs[wm]
 }
 
-func (wm *weakMap) remove(key *Object) bool {
-	ref := key.weakRef
-	if ref == nil {
-		return false
-	}
-	_, exists := wm.data[ref.id]
-	if exists {
-		delete(wm.data, ref.id)
-		key.runtime.removeWeakKey(ref.id, wm)
+func (wm weakMap) remove(key *Object) bool {
+	if _, exists := key.weakRefs[wm]; exists {
+		delete(key.weakRefs, wm)
+		return true
 	}
-	return exists
+	return false
 }
 
-func (wm *weakMap) has(key *Object) bool {
-	ref := key.weakRef
-	if ref == nil {
-		return false
-	}
-	_, exists := wm.data[ref.id]
+func (wm weakMap) has(key *Object) bool {
+	_, exists := key.weakRefs[wm]
 	return exists
 }
 

builtin_weakmap_test.go

@@ -1,33 +1,38 @@
 package goja
 
 import (
-	"runtime"
 	"testing"
 )
 
-func TestWeakMapExpiry(t *testing.T) {
+func TestWeakMap(t *testing.T) {
 	vm := New()
 	_, err := vm.RunString(`
 	var m = new WeakMap();
+	var m1 = new WeakMap();
 	var key = {};
 	m.set(key, true);
+	m1.set(key, false);
 	if (!m.has(key)) {
 		throw new Error("has");
 	}
 	if (m.get(key) !== true) {
 		throw new Error("value does not match");
 	}
-	key = undefined;
+	if (!m1.has(key)) {
+		throw new Error("has (m1)");
+	}
+	if (m1.get(key) !== false) {
+		throw new Error("m1 value does not match");
+	}
+	m.delete(key);
+	if (m.has(key)) {
+		throw new Error("m still has after delete");
+	}
+	if (!m1.has(key)) {
+		throw new Error("m1 does not have after delete from m");
+	}
 	`)
 	if err != nil {
 		t.Fatal(err)
 	}
-	runtime.GC()
-	runtime.GC()
-	vm.RunString("true") // this will trigger dead keys removal
-	wmo := vm.Get("m").ToObject(vm).self.(*weakMapObject)
-	l := len(wmo.m.data)
-	if l > 0 {
-		t.Fatal("Object has not been removed")
-	}
 }

builtin_weakset.go

@@ -1,55 +1,13 @@
 package goja
 
-type weakSet struct {
-	data map[uint64]struct{}
-}
-
 type weakSetObject struct {
 	baseObject
-	s *weakSet
-}
-
-func newWeakSet() *weakSet {
-	return &weakSet{
-		data: make(map[uint64]struct{}),
-	}
+	s weakMap
 }
 
 func (ws *weakSetObject) init() {
 	ws.baseObject.init()
-	ws.s = newWeakSet()
-}
-
-func (ws *weakSet) removeId(id uint64) {
-	delete(ws.data, id)
-}
-
-func (ws *weakSet) add(o *Object) {
-	ref := o.getWeakRef()
-	ws.data[ref.id] = struct{}{}
-	o.runtime.addWeakKey(ref.id, ws)
-}
-
-func (ws *weakSet) remove(o *Object) bool {
-	ref := o.weakRef
-	if ref == nil {
-		return false
-	}
-	_, exists := ws.data[ref.id]
-	if exists {
-		delete(ws.data, ref.id)
-		o.runtime.removeWeakKey(ref.id, ws)
-	}
-	return exists
-}
-
-func (ws *weakSet) has(o *Object) bool {
-	ref := o.weakRef
-	if ref == nil {
-		return false
-	}
-	_, exists := ws.data[ref.id]
-	return exists
+	ws.s = weakMap(ws.val.runtime.genId())
 }
 
 func (r *Runtime) weakSetProto_add(call FunctionCall) Value {
@@ -58,7 +16,7 @@ func (r *Runtime) weakSetProto_add(call FunctionCall) Value {
 	if !ok {
 		panic(r.NewTypeError("Method WeakSet.prototype.add called on incompatible receiver %s", thisObj.String()))
 	}
-	wso.s.add(r.toObject(call.Argument(0)))
+	wso.s.set(r.toObject(call.Argument(0)), nil)
 	return call.This
 }
 
@@ -119,7 +77,7 @@ func (r *Runtime) builtin_newWeakSet(args []Value, newTarget *Object) *Object {
 			if adder == r.global.weakSetAdder {
 				if arr := r.checkStdArrayIter(arg); arr != nil {
 					for _, v := range arr.values {
-						wso.s.add(r.toObject(v))
+						wso.s.set(r.toObject(v), nil)
 					}
 					return o
 				}

builtin_weakset_test.go

@@ -1,7 +1,6 @@
 package goja
 
 import (
-	"runtime"
 	"testing"
 )
 
@@ -21,30 +20,6 @@ func TestWeakSetBasic(t *testing.T) {
 	testScript1(SCRIPT, _undefined, t)
 }
 
-func TestWeakSetExpiry(t *testing.T) {
-	vm := New()
-	_, err := vm.RunString(`
-	var s = new WeakSet();
-	var o = {};
-	s.add(o);
-	if (!s.has(o)) {
-		throw new Error("has");
-	}
-	o = undefined;
-	`)
-	if err != nil {
-		t.Fatal(err)
-	}
-	runtime.GC()
-	runtime.GC()
-	vm.RunString("true") // this will trigger dead keys removal
-	wso := vm.Get("s").ToObject(vm).self.(*weakSetObject)
-	l := len(wso.s.data)
-	if l > 0 {
-		t.Fatal("Object has not been removed")
-	}
-}
-
 func TestWeakSetArraySimple(t *testing.T) {
 	const SCRIPT = `
 	var o1 = {}, o2 = {}, o3 = {};

object.go

@@ -4,9 +4,7 @@ import (
 	"fmt"
 	"math"
 	"reflect"
-	"runtime"
 	"sort"
-	"sync"
 
 	"github.com/dop251/goja/unistring"
 )
@@ -41,80 +39,12 @@ var (
 	hintString  Value = asciiString("string")
 )
 
-type weakCollection interface {
-	removeId(uint64)
-}
-
-type weakCollections struct {
-	objId uint64
-	colls []weakCollection
-}
-
-func (r *weakCollections) add(c weakCollection) {
-	for _, ec := range r.colls {
-		if ec == c {
-			return
-		}
-	}
-	r.colls = append(r.colls, c)
-}
-
-func (r *weakCollections) id() uint64 {
-	return r.objId
-}
-
-func (r *weakCollections) remove(c weakCollection) {
-	if cap(r.colls) > 16 && cap(r.colls)>>2 > len(r.colls) {
-		// shrink
-		colls := make([]weakCollection, 0, len(r.colls))
-		for _, coll := range r.colls {
-			if coll != c {
-				colls = append(colls, coll)
-			}
-		}
-		r.colls = colls
-	} else {
-		for i, coll := range r.colls {
-			if coll == c {
-				l := len(r.colls) - 1
-				r.colls[i] = r.colls[l]
-				r.colls[l] = nil
-				r.colls = r.colls[:l]
-				break
-			}
-		}
-	}
-}
-
-func finalizeObjectWeakRefs(r *objectWeakRef) {
-	r.tracker.add(r.id)
-}
-
-type weakRefTracker struct {
-	sync.Mutex
-	list []uint64
-}
-
-func (t *weakRefTracker) add(id uint64) {
-	t.Lock()
-	t.list = append(t.list, id)
-	t.Unlock()
-}
-
-// An object that gets finalized when the corresponding *Object is garbage-collected.
-// It must be ensured that neither the *Object, nor the *Runtime is reachable from this struct,
-// otherwise it will create a circular reference with a Finalizer which will make it not garbage-collectable.
-type objectWeakRef struct {
-	id      uint64
-	tracker *weakRefTracker
-}
-
 type Object struct {
 	id      uint64
 	runtime *Runtime
 	self    objectImpl
 
-	weakRef *objectWeakRef
+	weakRefs map[weakMap]Value
 }
 
 type iterNextFunc func() (propIterItem, iterNextFunc)
@@ -1484,31 +1414,22 @@ func (o *Object) defineOwnProperty(n Value, desc PropertyDescriptor, throw bool)
 	}
 }
 
-func (o *Object) getWeakRef() *objectWeakRef {
-	if o.weakRef == nil {
-		if o.runtime.weakRefTracker == nil {
-			o.runtime.weakRefTracker = &weakRefTracker{}
-		}
-		o.weakRef = &objectWeakRef{
-			id:      o.getId(),
-			tracker: o.runtime.weakRefTracker,
-		}
-		runtime.SetFinalizer(o.weakRef, finalizeObjectWeakRefs)
+func (o *Object) getWeakRefs() map[weakMap]Value {
+	refs := o.weakRefs
+	if refs == nil {
+		refs = make(map[weakMap]Value)
+		o.weakRefs = refs
 	}
-
-	return o.weakRef
+	return refs
 }
 
 func (o *Object) getId() uint64 {
-	for o.id == 0 {
-		if o.runtime.hash == nil {
-			h := o.runtime.getHash()
-			o.runtime.idSeq = h.Sum64()
-		}
-		o.id = o.runtime.idSeq
-		o.runtime.idSeq++
+	id := o.id
+	if id == 0 {
+		id = o.runtime.genId()
+		o.id = id
 	}
-	return o.id
+	return id
 }
 
 func (o *guardedObject) guard(props ...unistring.String) {

runtime.go

@@ -172,16 +172,6 @@ type Runtime struct {
 	vm    *vm
 	hash  *maphash.Hash
 	idSeq uint64
-
-	// Contains a list of ids of finalized weak keys so that the runtime could pick it up and remove from
-	// all weak collections using the weakKeys map. The runtime picks it up either when the topmost function
-	// returns (i.e. the callstack becomes empty) or every 10000 'ticks' (vm instructions).
-	// It is implemented this way to avoid circular references which at the time of writing (go 1.15) causes
-	// the whole structure to become not garbage-collectable.
-	weakRefTracker *weakRefTracker
-
-	// Contains a list of weak collections that contain the key with the id.
-	weakKeys map[uint64]*weakCollections
 }
 
 type StackFrame struct {
@@ -2267,57 +2257,9 @@ func (r *Runtime) getHash() *maphash.Hash {
 	return r.hash
 }
 
-func (r *Runtime) addWeakKey(id uint64, coll weakCollection) {
-	keys := r.weakKeys
-	if keys == nil {
-		keys = make(map[uint64]*weakCollections)
-		r.weakKeys = keys
-	}
-	colls := keys[id]
-	if colls == nil {
-		colls = &weakCollections{
-			objId: id,
-		}
-		keys[id] = colls
-	}
-	colls.add(coll)
-}
-
-func (r *Runtime) removeWeakKey(id uint64, coll weakCollection) {
-	keys := r.weakKeys
-	if colls := keys[id]; colls != nil {
-		colls.remove(coll)
-		if len(colls.colls) == 0 {
-			delete(keys, id)
-		}
-	}
-}
-
-// this gets inlined so a CALL is avoided on a critical path
-func (r *Runtime) removeDeadKeys() {
-	if r.weakRefTracker != nil {
-		r.doRemoveDeadKeys()
-	}
-}
-
-func (r *Runtime) doRemoveDeadKeys() {
-	r.weakRefTracker.Lock()
-	list := r.weakRefTracker.list
-	r.weakRefTracker.list = nil
-	r.weakRefTracker.Unlock()
-	for _, id := range list {
-		if colls := r.weakKeys[id]; colls != nil {
-			for _, coll := range colls.colls {
-				coll.removeId(id)
-			}
-			delete(r.weakKeys, id)
-		}
-	}
-}
-
 // called when the top level function returns (i.e. control is passed outside the Runtime).
 func (r *Runtime) leave() {
-	r.removeDeadKeys()
+	// run jobs, etc...
 }
 
 func nilSafe(v Value) Value {
@@ -2396,3 +2338,16 @@ func growCap(newSize, oldSize, oldCap int) int {
 		}
 	}
 }
+
+func (r *Runtime) genId() (ret uint64) {
+	if r.hash == nil {
+		h := r.getHash()
+		r.idSeq = h.Sum64()
+	}
+	if r.idSeq == 0 {
+		r.idSeq = 1
+	}
+	ret = r.idSeq
+	r.idSeq++
+	return
+}

vm.go

@@ -308,7 +308,6 @@ func (vm *vm) run() {
 		ticks++
 		if ticks > 10000 {
 			runtime.Gosched()
-			vm.r.removeDeadKeys()
 			ticks = 0
 		}
 	}