From 91f791400338673f025897c70da99301c85694c5 Mon Sep 17 00:00:00 2001
From: "Miss Islington (bot)"
 <31488909+miss-islington@users.noreply.github.com>
Date: Thu, 14 Jul 2022 09:47:17 -0700
Subject: GH-89988: Fix memory leak in pickle.Pickler dispatch_table lookup
 (GH-94298)

(cherry picked from commit 01ef1f95dab9c9930ce1a23634a3e5a8331bf3c7)

Co-authored-by: Kumar Aditya <59607654+kumaraditya303@users.noreply.github.com>
---
 Lib/test/test_pickle.py                            | 23 ++++++++++++++++++++++
 .../2022-06-26-10-59-15.gh-issue-89988.K8rnmt.rst  |  1 +
 Modules/_pickle.c                                  |  4 +++-
 3 files changed, 27 insertions(+), 1 deletion(-)
 create mode 100644 Misc/NEWS.d/next/Library/2022-06-26-10-59-15.gh-issue-89988.K8rnmt.rst

diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py
index 057af21..44fdca7 100644
--- a/Lib/test/test_pickle.py
+++ b/Lib/test/test_pickle.py
@@ -155,6 +155,29 @@ class PyIdPersPicklerTests(AbstractIdentityPersistentPicklerTests,
         check(PersPickler)
 
     @support.cpython_only
+    def test_custom_pickler_dispatch_table_memleak(self):
+        # See https://github.com/python/cpython/issues/89988
+
+        class Pickler(self.pickler):
+            def __init__(self, *args, **kwargs):
+                self.dispatch_table = table
+                super().__init__(*args, **kwargs)
+
+        class DispatchTable:
+            pass
+
+        table = DispatchTable()
+        pickler = Pickler(io.BytesIO())
+        self.assertIs(pickler.dispatch_table, table)
+        table_ref = weakref.ref(table)
+        self.assertIsNotNone(table_ref())
+        del pickler
+        del table
+        support.gc_collect()
+        self.assertIsNone(table_ref())
+
+
+    @support.cpython_only
     def test_unpickler_reference_cycle(self):
         def check(Unpickler):
             for proto in range(pickle.HIGHEST_PROTOCOL + 1):
diff --git a/Misc/NEWS.d/next/Library/2022-06-26-10-59-15.gh-issue-89988.K8rnmt.rst b/Misc/NEWS.d/next/Library/2022-06-26-10-59-15.gh-issue-89988.K8rnmt.rst
new file mode 100644
index 0000000..811a8d6
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2022-06-26-10-59-15.gh-issue-89988.K8rnmt.rst
@@ -0,0 +1 @@
+Fix memory leak in :class:`pickle.Pickler` when looking up :attr:`dispatch_table`. Patch by Kumar Aditya.
diff --git a/Modules/_pickle.c b/Modules/_pickle.c
index 23d26f6..b4b1dda 100644
--- a/Modules/_pickle.c
+++ b/Modules/_pickle.c
@@ -4773,7 +4773,9 @@ _pickle_Pickler___init___impl(PicklerObject *self, PyObject *file,
     {
         return -1;
     }
-
+    if (self->dispatch_table != NULL) {
+        return 0;
+    }
     if (_PyObject_LookupAttr((PyObject *)self, &_Py_ID(dispatch_table),
                              &self->dispatch_table) < 0) {
         return -1;
-- 
cgit v0.12