From a896a0121934cad3d3fe2dd859fc16898666c811 Mon Sep 17 00:00:00 2001
From: Thomas Grainger <tagrain@gmail.com>
Date: Sat, 2 Apr 2022 11:50:59 +0100
Subject: [PATCH] shutdown all profile threads before running
 test_weakref_cache

---
 distributed/tests/test_spill.py | 81 ++++++++++++++++++---------------
 1 file changed, 45 insertions(+), 36 deletions(-)

diff --git a/distributed/tests/test_spill.py b/distributed/tests/test_spill.py
index 0b185a8239..e2c91fb006 100644
--- a/distributed/tests/test_spill.py
+++ b/distributed/tests/test_spill.py
@@ -332,6 +332,7 @@ def _print_chain(o_wr):
     return o
 
 
+@pytest.mark.parametrize("attempt", range(1000))
 @pytest.mark.parametrize(
     "cls,expect_cached",
     [
@@ -340,39 +341,47 @@ def _print_chain(o_wr):
     ],
 )
 @pytest.mark.parametrize("size", [60, 110])
-def test_weakref_cache(tmpdir, cls, expect_cached, size):
-    buf = SpillBuffer(str(tmpdir), target=100)
-
-    # Run this test twice:
-    # - x is smaller than target and is evicted by y;
-    # - x is individually larger than target and it never touches fast
-    x = cls(size)
-    canary = weakref.ref(x.canary)
-    buf["x"] = x
-    if size < 100:
-        buf["y"] = cls(60)  # spill x
-    assert "x" in buf.slow
-
-    # Test that we update the weakref cache on setitem
-    assert (buf["x"] is x) == expect_cached
-
-    # Do not use id_x = id(x), as in CPython id's are C memory addresses and are reused
-    # by PyMalloc when you descope objects, so a brand new object might end up having
-    # the same id as a deleted one
-    id_x = x.id
-    del x
-
-    if size < 100:
-        buf["y"]
-
-    assert _print_chain(canary) is None
-    assert "x" in buf.slow
-
-    x2 = buf["x"]
-    assert x2.id != id_x
-    if size < 100:
-        buf["y"]
-    assert "x" in buf.slow
-
-    # Test that we update the weakref cache on getitem
-    assert (buf["x"] is x2) == expect_cached
+def test_weakref_cache(tmpdir, cls, expect_cached, size, attempt):
+    from distributed.profile import shutdown_profile_threads
+
+    shutdown_profile_threads()
+
+    gc.disable()
+    try:
+        buf = SpillBuffer(str(tmpdir), target=100)
+
+        # Run this test twice:
+        # - x is smaller than target and is evicted by y;
+        # - x is individually larger than target and it never touches fast
+        x = cls(size)
+        canary = weakref.ref(x.canary)
+        buf["x"] = x
+        if size < 100:
+            buf["y"] = cls(60)  # spill x
+        assert "x" in buf.slow
+
+        # Test that we update the weakref cache on setitem
+        assert (buf["x"] is x) == expect_cached
+
+        # Do not use id_x = id(x), as in CPython id's are C memory addresses and are reused
+        # by PyMalloc when you descope objects, so a brand new object might end up having
+        # the same id as a deleted one
+        id_x = x.id
+        del x
+
+        if size < 100:
+            buf["y"]
+
+        assert canary() is None
+        assert "x" in buf.slow
+
+        x2 = buf["x"]
+        assert x2.id != id_x
+        if size < 100:
+            buf["y"]
+        assert "x" in buf.slow
+
+        # Test that we update the weakref cache on getitem
+        assert (buf["x"] is x2) == expect_cached
+    finally:
+        gc.enable()