Bug 680937, part 2 - Add native wrapper preservation hook, call it in WeakMap_set. r=billm

js/src/jsapi.cpp

@@ -723,6 +723,7 @@ JSRuntime::JSRuntime()
     trustedPrincipals_(NULL),
     wrapObjectCallback(NULL),
     preWrapObjectCallback(NULL),
+    preserveWrapperCallback(NULL),
     inOOMReport(0)
 {
     /* Initialize infallibly first, so we can goto bad and JS_DestroyRuntime. */

js/src/jscntxt.h

@@ -650,6 +650,7 @@ struct JSRuntime
 
     JSWrapObjectCallback wrapObjectCallback;
     JSPreWrapCallback    preWrapObjectCallback;
+    js::PreserveWrapperCallback preserveWrapperCallback;
 
     /*
      * To ensure that cx->malloc does not cause a GC, we set this flag during

js/src/jsfriendapi.cpp

@@ -305,6 +305,12 @@ js::SetFunctionNativeReserved(JSObject *fun, size_t which, const Value &val)
     fun->toFunction()->setExtendedSlot(which, val);
 }
 
+void
+js::SetPreserveWrapperCallback(JSRuntime *rt, PreserveWrapperCallback callback)
+{
+    rt->preserveWrapperCallback = callback;
+}
+
 /*
  * The below code is for temporary telemetry use. It can be removed when
  * sufficient data has been harvested.

js/src/jsfriendapi.h

@@ -168,6 +168,9 @@ JS_END_EXTERN_C
 
 namespace js {
 
+typedef bool
+(* PreserveWrapperCallback)(JSContext *cx, JSObject *obj);
+
 #ifdef DEBUG
  /*
   * DEBUG-only method to dump the complete object graph of heap-allocated things.
@@ -425,6 +428,9 @@ GetPropertyNames(JSContext *cx, JSObject *obj, uintN flags, js::AutoIdVector *pr
 JS_FRIEND_API(bool)
 StringIsArrayIndex(JSLinearString *str, jsuint *indexp);
 
+JS_FRIEND_API(void)
+SetPreserveWrapperCallback(JSRuntime *rt, PreserveWrapperCallback callback);
+
 /*
  * NB: these flag bits are encoded into the bytecode stream in the immediate
  * operand of JSOP_ITER, so don't change them without advancing jsxdrapi.h's

js/src/jsweakmap.cpp

@@ -233,6 +233,15 @@ WeakMap_set(JSContext *cx, uintN argc, Value *vp)
 
     if (!map->put(key, value))
         goto out_of_memory;
+
+    // Preserve wrapped native keys to prevent wrapper optimization.
+    if (key->getClass()->ext.isWrappedNative) {
+        if (!cx->runtime->preserveWrapperCallback ||
+            !cx->runtime->preserveWrapperCallback(cx, key)) {
+            JS_ReportWarning(cx, "Failed to preserve wrapper of wrapped native weak map key.");
+        }
+    }
+
     args.rval().setUndefined();
     return true;
 

js/xpconnect/tests/chrome/test_weakmaps.xul

@@ -61,26 +61,6 @@ https://bugzilla.mozilla.org/show_bug.cgi?id=668855
   let weakref = make_gray_loop();
 
 
-  /* Weak map entries where the key is a dead XPCWrappedNative key should be removed.
-
-     map2 is a weak map that is black, so its entries will be visited during the 
-     black marking phase.  We then use a new div element as key.  The div element is dead 
-     after dead_xpc_key returns, so the weak map entry should be removed.
-
-     The simple wrapper deoptimization for XPCWrappedNative weak map keys implemented in
-     Bug 655297 will end up marking the key black, keeping the entry from being collected.
-
-  */
-  let map2 = new WeakMap;
-
-  let dead_xpc_key = function () {
-    let div = document.createElement("div");
-    map2.set(div, 0);
-  };
-
-  dead_xpc_key();
-
-
   /* Combinations of live and dead gray maps/keys. */
   let basic_weak_ref = null;
   let basic_map_weak_ref = null;
@@ -203,6 +183,37 @@ https://bugzilla.mozilla.org/show_bug.cgi?id=668855
   chained_garbage_maps();
 
 
+  /* black weak map, chained garbage cycle involving DOM, XPCWN keys */
+  let wn_garbage_map = new WeakMap;
+
+  let wn_chained_garbage_maps = function () {
+    let dom0 = document.createElement("div");
+    let dom = dom0;
+    for(let i = 0; i < num_chains; i++) {
+      let new_dom = document.createElement("div");
+      wn_garbage_map.set(dom, {wn_val_child:new_dom});
+      dom = document.createElement("div");
+      new_dom.appendChild(dom);
+    };
+    // tie the knot
+    dom.appendChild(dom0);
+  };
+
+  wn_chained_garbage_maps();
+
+
+  /* The cycle collector shouldn't remove a live wrapped native key. */
+
+  let wn_live_map = new WeakMap;
+
+  let make_live_map = function () {
+    let live = get_live_dom();
+    wn_live_map.set(live, {});
+  }
+
+  make_live_map();
+
+
   /* set up for running precise GC/CC then checking the results */
 
   SimpleTest.waitForExplicitFinish();
@@ -220,15 +231,20 @@ https://bugzilla.mozilla.org/show_bug.cgi?id=668855
 
     ok(weak_ref_dead(weakref), "Garbage gray cycle should be collected.");
 
-    // this will fail without the XPCwrapped native key fix (Bug 680937)
-    todo_is(Cu.nondeterministicGetWeakMapKeys(map2).length, 0, "Dead XPCWrappedNative keys should be collected.");
-
     check_nested_cc_maps();
 
     is(Cu.nondeterministicGetWeakMapKeys(garbage_map).length, 0, "Chained garbage weak map entries should not leak.");
 
     check_basic_unit();
 
+    // fixed by Bug 680937
+    is(Cu.nondeterministicGetWeakMapKeys(wn_garbage_map).length, 0,
+       "Chained garbage WN weak map entries should not leak.");
+
+    // fixed by Bug 680937
+    is(Cu.nondeterministicGetWeakMapKeys(wn_live_map).length, 1,
+       "Live weak map wrapped native key should not be removed.");
+
     SimpleTest.finish();
   });