Rewrote the handling of weak keys to avoid creating circular structures when weak keys are reachable from Runtime. Fixes #199.

Author: dop251

README.md

@@ -21,6 +21,24 @@ Features
  * Sourcemaps.
  * Some ES6 functionality, still work in progress, see https://github.com/dop251/goja/milestone/1?closed=1
 
+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:
+
+```go
+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.
+```
+
+Note, this does not have any effect on the application logic, but causes a higher-than-expected memory usage.
+
 FAQ
 ---
 

builtin_weakmap.go

@@ -1,11 +1,6 @@
 package goja
 
-import "sync"
-
 type weakMap struct {
-	// need to synchronise access to the data map because it may be accessed
-	// from the finalizer goroutine
-	sync.Mutex
 	data map[uint64]Value
 }
 
@@ -26,57 +21,43 @@ func (wmo *weakMapObject) init() {
 }
 
 func (wm *weakMap) removeId(id uint64) {
-	wm.Lock()
 	delete(wm.data, id)
-	wm.Unlock()
 }
 
 func (wm *weakMap) set(key *Object, value Value) {
-	refs := key.getWeakCollRefs()
-	wm.Lock()
-	wm.data[refs.id()] = value
-	wm.Unlock()
-	refs.add(wm)
+	ref := key.getWeakRef()
+	wm.data[ref.id] = value
+	key.runtime.addWeakKey(ref.id, wm)
 }
 
 func (wm *weakMap) get(key *Object) Value {
-	refs := key.weakColls
-	if refs == nil {
+	ref := key.weakRef
+	if ref == nil {
 		return nil
 	}
-	wm.Lock()
-	ret := wm.data[refs.id()]
-	wm.Unlock()
+	ret := wm.data[ref.id]
 	return ret
 }
 
 func (wm *weakMap) remove(key *Object) bool {
-	refs := key.weakColls
-	if refs == nil {
+	ref := key.weakRef
+	if ref == nil {
 		return false
 	}
-	id := refs.id()
-	wm.Lock()
-	_, exists := wm.data[id]
-	if exists {
-		delete(wm.data, id)
-	}
-	wm.Unlock()
+	_, exists := wm.data[ref.id]
 	if exists {
-		refs.remove(wm)
+		delete(wm.data, ref.id)
+		key.runtime.removeWeakKey(ref.id, wm)
 	}
 	return exists
 }
 
 func (wm *weakMap) has(key *Object) bool {
-	refs := key.weakColls
-	if refs == nil {
+	ref := key.weakRef
+	if ref == nil {
 		return false
 	}
-	id := refs.id()
-	wm.Lock()
-	_, exists := wm.data[id]
-	wm.Unlock()
+	_, exists := wm.data[ref.id]
 	return exists
 }
 

builtin_weakmap_test.go

@@ -24,10 +24,9 @@ func TestWeakMapExpiry(t *testing.T) {
 	}
 	runtime.GC()
 	runtime.GC()
+	vm.RunString("true") // this will trigger dead keys removal
 	wmo := vm.Get("m").ToObject(vm).self.(*weakMapObject)
-	wmo.m.Lock()
 	l := len(wmo.m.data)
-	wmo.m.Unlock()
 	if l > 0 {
 		t.Fatal("Object has not been removed")
 	}

builtin_weakset.go

@@ -1,11 +1,6 @@
 package goja
 
-import "sync"
-
 type weakSet struct {
-	// need to synchronise access to the data map because it may be accessed
-	// from the finalizer goroutine
-	sync.Mutex
 	data map[uint64]struct{}
 }
 
@@ -26,43 +21,34 @@ func (ws *weakSetObject) init() {
 }
 
 func (ws *weakSet) removeId(id uint64) {
-	ws.Lock()
 	delete(ws.data, id)
-	ws.Unlock()
 }
 
 func (ws *weakSet) add(o *Object) {
-	refs := o.getWeakCollRefs()
-	ws.Lock()
-	ws.data[refs.id()] = struct{}{}
-	ws.Unlock()
-	refs.add(ws)
+	ref := o.getWeakRef()
+	ws.data[ref.id] = struct{}{}
+	o.runtime.addWeakKey(ref.id, ws)
 }
 
 func (ws *weakSet) remove(o *Object) bool {
-	if o.weakColls == nil {
+	ref := o.weakRef
+	if ref == nil {
 		return false
 	}
-	id := o.weakColls.id()
-	ws.Lock()
-	_, exists := ws.data[id]
-	if exists {
-		delete(ws.data, id)
-	}
-	ws.Unlock()
+	_, exists := ws.data[ref.id]
 	if exists {
-		o.weakColls.remove(ws)
+		delete(ws.data, ref.id)
+		o.runtime.removeWeakKey(ref.id, ws)
 	}
 	return exists
 }
 
 func (ws *weakSet) has(o *Object) bool {
-	if o.weakColls == nil {
+	ref := o.weakRef
+	if ref == nil {
 		return false
 	}
-	ws.Lock()
-	_, exists := ws.data[o.weakColls.id()]
-	ws.Unlock()
+	_, exists := ws.data[ref.id]
 	return exists
 }
 

builtin_weakset_test.go

@@ -37,10 +37,9 @@ func TestWeakSetExpiry(t *testing.T) {
 	}
 	runtime.GC()
 	runtime.GC()
+	vm.RunString("true") // this will trigger dead keys removal
 	wso := vm.Get("s").ToObject(vm).self.(*weakSetObject)
-	wso.s.Lock()
 	l := len(wso.s.data)
-	wso.s.Unlock()
 	if l > 0 {
 		t.Fatal("Object has not been removed")
 	}

object.go

@@ -6,6 +6,7 @@ import (
 	"reflect"
 	"runtime"
 	"sort"
+	"sync"
 
 	"github.com/dop251/goja/unistring"
 )
@@ -84,25 +85,35 @@ func (r *weakCollections) remove(c weakCollection) {
 	}
 }
 
-func finalizeObjectWeakRefs(r *weakCollections) {
-	id := r.id()
-	for _, c := range r.colls {
-		c.removeId(id)
-	}
-	r.colls = nil
+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
 
-	// Contains references to all weak collections that contain this Object.
-	// weakColls has a finalizer that removes the Object's id from all weak collections.
-	// The id is the weakColls pointer value converted to uintptr.
-	// Note, cannot set the finalizer on the *Object itself because it's a part of a
-	// reference cycle.
-	weakColls *weakCollections
+	weakRef *objectWeakRef
 }
 
 type iterNextFunc func() (propIterItem, iterNextFunc)
@@ -1399,15 +1410,19 @@ func (o *Object) defineOwnProperty(n Value, desc PropertyDescriptor, throw bool)
 	}
 }
 
-func (o *Object) getWeakCollRefs() *weakCollections {
-	if o.weakColls == nil {
-		o.weakColls = &weakCollections{
-			objId: o.getId(),
+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.weakColls, finalizeObjectWeakRefs)
+		runtime.SetFinalizer(o.weakRef, finalizeObjectWeakRefs)
 	}
 
-	return o.weakColls
+	return o.weakRef
 }
 
 func (o *Object) getId() uint64 {

runtime.go

@@ -170,6 +170,16 @@ 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 {
@@ -1195,6 +1205,7 @@ func (r *Runtime) RunProgram(p *Program) (result Value, err error) {
 		r.vm.clearStack()
 	} else {
 		r.vm.stack = nil
+		r.leave()
 	}
 	return
 }
@@ -1966,7 +1977,11 @@ func AssertFunction(v Value) (Callable, bool) {
 				if ex != nil {
 					err = ex
 				}
-				obj.runtime.vm.clearStack()
+				vm := obj.runtime.vm
+				vm.clearStack()
+				if len(vm.callStack) == 0 {
+					obj.runtime.leave()
+				}
 				return
 			}, true
 		}
@@ -2184,6 +2199,59 @@ 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()
+}
+
 func nilSafe(v Value) Value {
 	if v != nil {
 		return v

vm.go

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