From a3957d6aafca2985ad7a15b70a9cc0dc563f5a0f Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sat, 8 Sep 2018 19:32:20 +0100
Subject: [PATCH] parent: add Context.forget_chain().

---
 docs/api.rst                | 11 ++++++++---
 mitogen/core.py             |  5 +++++
 mitogen/parent.py           |  3 +++
 tests/call_function_test.py |  8 ++++++++
 4 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/docs/api.rst b/docs/api.rst
index 395147b2..56bb5c1f 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -934,9 +934,9 @@ Context Class
                 if recv.get().unpickle() == 'baz':
                     pass
 
-            Note that for long-lived programs, there is presently no mechanism
-            for clearing the chain history on a target. This will be addressed
-            in future.
+            It is necessary to explicitly clean up the chain history on a
+            target, otherwise unbounded memory usage is possible. See
+            :meth:`forget_chain`.
 
         :returns:
             :class:`mitogen.core.Receiver` configured to receive the result
@@ -977,6 +977,11 @@ Context Class
         :raises mitogen.core.CallError:
             An exception was raised in the remote context during execution.
 
+    .. method:: forget_chain (chain_id)
+
+        Instruct the target to forget any exception related to `chain_id`, a
+        key previously used as the `mitogen_chain` parameter to
+        :meth:`call_async`.
 
 
 Receiver Class
diff --git a/mitogen/core.py b/mitogen/core.py
index c959ea5d..0142be7c 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -1959,6 +1959,11 @@ class Dispatcher(object):
                              policy=has_parent_authority)
         listen(econtext.broker, 'shutdown', self.recv.close)
 
+    @classmethod
+    @takes_econtext
+    def forget_chain(cls, chain_id, econtext):
+        econtext.dispatcher._error_by_chain_id.pop(chain_id, None)
+
     def _parse_request(self, msg):
         data = msg.unpickle(throw=False)
         _v and LOG.debug('_dispatch_one(%r)', data)
diff --git a/mitogen/parent.py b/mitogen/parent.py
index d7177dde..5c7fb642 100644
--- a/mitogen/parent.py
+++ b/mitogen/parent.py
@@ -1160,6 +1160,9 @@ class Context(mitogen.core.Context):
                   self, fn, args, kwargs)
         self.send(make_call_msg(fn, *args, **kwargs))
 
+    def forget_chain(self, chain_id):
+        self.call_no_reply(mitogen.core.Dispatcher.forget_chain, chain_id)
+
     def shutdown(self, wait=False):
         LOG.debug('%r.shutdown() sending SHUTDOWN', self)
         latch = mitogen.core.Latch()
diff --git a/tests/call_function_test.py b/tests/call_function_test.py
index b8b07283..ca56f07a 100644
--- a/tests/call_function_test.py
+++ b/tests/call_function_test.py
@@ -143,6 +143,14 @@ class ChainTest(testlib.RouterMixin, testlib.TestCase):
             lambda: self.local.call(func_returns_arg, 'yes', mitogen_chain='c1'))
         self.local.call_no_reply(function_that_fails, 'c2', mitogen_chain='c2')
 
+    def test_forget(self):
+        self.local.call_no_reply(function_that_fails, 'x1', mitogen_chain='c1')
+        e1 = self.assertRaises(mitogen.core.CallError,
+            lambda: self.local.call(function_that_fails, 'x2', mitogen_chain='c1'))
+        self.local.forget_chain('c1')
+        self.assertEquals('x3',
+            self.local.call(func_returns_arg, 'x3', mitogen_chain='c1'))
+
 
 if __name__ == '__main__':
     unittest2.main()