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

Author: dop251

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 = {};