From bd8f2f936b53d30e61a6fe04a90cc933cc961b93 Mon Sep 17 00:00:00 2001
From: Emmanuel Lepage Vallee <elv1313@gmail.com>
Date: Sun, 19 Sep 2021 22:21:38 -0700
Subject: [PATCH] naughty: Fix a memory leak related to suspended
 notifications.

---
 lib/naughty/core.lua          | 15 ++++++++++----
 tests/test-naughty-screen.lua | 38 ++++++++++++++++++++++++++++++++---
 2 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/lib/naughty/core.lua b/lib/naughty/core.lua
index 0afd55083..757bbe172 100644
--- a/lib/naughty/core.lua
+++ b/lib/naughty/core.lua
@@ -191,6 +191,10 @@ naughty.notifications = { suspended = { }, _expired = {{}} }
 
 naughty._active = {}
 
+local function get_screen(s)
+    return s and capi.screen[s]
+end
+
 local function init_screen(s)
     if naughty.notifications[s] then return end
 
@@ -218,15 +222,18 @@ capi.screen.connect_signal("removed", function(scr)
             naughty.emit_signal("request::screen", n, "removed", {})
         end
     end
+
+    for _, n in ipairs(naughty._active) do
+        if n._private.args and get_screen(n._private.args.screen) == scr then
+            n._private.args.screen = nil
+        end
+    end
+
     -- Destroy all notifications on this screen
     naughty.destroy_all_notifications({scr})
     naughty.notifications[scr] = nil
 end)
 
-local function get_screen(s)
-    return s and capi.screen[s]
-end
-
 local function remove_from_index(n)
     for _, positions in pairs(naughty.notifications) do
         for _, ns in pairs(positions) do
diff --git a/tests/test-naughty-screen.lua b/tests/test-naughty-screen.lua
index 7428a462f..ef95acc41 100644
--- a/tests/test-naughty-screen.lua
+++ b/tests/test-naughty-screen.lua
@@ -30,11 +30,34 @@ local s1, s2 = mouse.screen, nil
 
 for _, p in ipairs(positions) do
     objs[p] = setmetatable({},{
-        __index=function(t,k) t[k]={};return t[k] end,
+        __index = function(t,k)
+            t[k] = setmetatable({}, {__mode = "kv"})
+            return t[k]
+        end,
         __mode = "k"
     })
 end
 
+local function cleanup(n)
+    -- Wait until there is no notifications left.
+    for _, pos in ipairs(positions) do
+        for s, notifs in pairs(objs[pos]) do
+            for k, n2 in ipairs(notifs) do
+                if n == n2 then
+                    table.remove(notifs, k)
+                    if #notifs == 0 then
+                        objs[pos][s] = nil
+                    end
+                    return
+                end
+            end
+        end
+    end
+end
+
+naughty.connect_signal("property::screen", cleanup)
+naughty.connect_signal("destroyed", cleanup)
+
 local function add_many(s)
     for _, pos in ipairs(positions) do
         for i=1, 5 do
@@ -143,9 +166,17 @@ for _, legacy_preset in ipairs {true, false} do
     table.insert(steps, function()
         s2:fake_remove()
 
-        -- Help the weak tables a little.
+        return true
+    end)
+
+    -- Check if notifications are moved.
+    table.insert(steps, function()
+        -- Wait until there is no notifications left.
         for _, pos in ipairs(positions) do
-            objs[pos][s1] = nil
+            if #objs[pos][s2] > 0 then
+                collectgarbage("collect")
+                return
+            end
         end
 
         -- Drop our string reference to s2.
@@ -154,6 +185,7 @@ for _, legacy_preset in ipairs {true, false} do
         return true
     end)
 
+
     table.insert(steps, function()
         if weak[1] == nil then return true end