From 9f94fb78c858cda6defb3fec1d1e72f108610873 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Fri, 13 Apr 2018 15:16:54 +0100
Subject: [PATCH] issue #159: make LRU size configurable.

---
 ansible_mitogen/services.py                   |  7 +++--
 docs/ansible.rst                              | 28 ++++++++++---------
 .../context_service/lru_one_target.yml        | 12 ++++++--
 tests/ansible/run_ansible_playbook.sh         |  1 +
 4 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/ansible_mitogen/services.py b/ansible_mitogen/services.py
index cfccd6ff..ba8d5eeb 100644
--- a/ansible_mitogen/services.py
+++ b/ansible_mitogen/services.py
@@ -39,9 +39,10 @@ when a child has completed a job.
 
 from __future__ import absolute_import
 import logging
-import sys
+import os
 import os.path
 import pprint
+import sys
 import threading
 import zlib
 
@@ -73,7 +74,7 @@ class ContextService(mitogen.service.Service):
     """
     handle = 500
     max_message_size = 1000
-    max_contexts = 20
+    max_interpreters = int(os.getenv('MITOGEN_MAX_INTERPRETERS', '20'))
 
     def __init__(self, *args, **kwargs):
         super(ContextService, self).__init__(*args, **kwargs)
@@ -127,7 +128,7 @@ class ContextService(mitogen.service.Service):
             return
 
         lru = self._lru_by_via.setdefault(via, [])
-        if len(lru) < self.max_contexts:
+        if len(lru) < self.max_interpreters:
             lru.append(new_context)
             return
 
diff --git a/docs/ansible.rst b/docs/ansible.rst
index e3b43163..c1199cf5 100644
--- a/docs/ansible.rst
+++ b/docs/ansible.rst
@@ -369,11 +369,10 @@ future.
 Interpreter Recycling
 ~~~~~~~~~~~~~~~~~~~~~
 
-To prevent accidental DoS, the extension stops creating persistent interpreters
-after the 20th interpreter has been created. Instead the most recently created
-interpreter is shut down to make room for any new interpreter. This is to avoid
-situations like below from triggering memory exhaustion by spawning a huge
-number of interpreters.
+The extension stops limits the number of persistent interpreters in use. When
+the limit is reached, the youngest interpreter is terminated before starting a
+new interpreter, preventing situations like below from triggering memory
+exhaustion.
 
 .. code-block:: yaml
 
@@ -392,15 +391,18 @@ number of interpreters.
             dest: "~{{item}}/.bashrc"
           with_items: "{{user_directory}}"
 
-The recycling behaviour does not occur for direct connections from the
-controller, and it is keyed on a per-host basis, i.e. up to 20 interpreters may
-exist for each directly connected target host.
+This recycling does not occur for direct connections from the controller, and
+it is keyed on a per-target basis, i.e. up to 20 interpreters may exist for
+each directly connected target.
 
-The newest interpreter is chosen to avoid recycling useful accounts, like
-"root" or "postgresql" that tend to appear early in a run, however it is simple
-to construct a playbook that defeats this strategy. A future version will key
-interpreters on the identity of the task, file and/or playbook that created
-them, avoiding the recycling of useful accounts in every scenario.
+The youngest interpreter is chosen to preserve useful accounts, like "root" or
+"postgresql" that tend to appear early in a run, however it is simple to
+construct a playbook that defeats this strategy. A future version will key
+interpreters on the identity of their creating task, file and/or playbook,
+avoiding useful account recycling in every scenario.
+
+To raise or lower the limit from 20, set the ``MITOGEN_MAX_INTERPRETERS``
+environment variable to a new value.
 
 
 Runtime Patches
diff --git a/tests/ansible/integration/context_service/lru_one_target.yml b/tests/ansible/integration/context_service/lru_one_target.yml
index 16fd7486..d4f9f8ad 100644
--- a/tests/ansible/integration/context_service/lru_one_target.yml
+++ b/tests/ansible/integration/context_service/lru_one_target.yml
@@ -2,6 +2,10 @@
 
 - hosts: all
   any_errors_fatal: true
+  vars:
+    max_interps: "{{lookup('env', 'MITOGEN_MAX_INTERPRETERS')}}"
+    ubound: "{{max_interps|int + 1}}"
+
   tasks:
     - name: integration/context_service/lru_one_target.yml
       assert:
@@ -12,7 +16,7 @@
       become: true
       vars:
         ansible_become_user: "mitogen__user{{item}}"
-      with_sequence: start=1 end=21
+      with_sequence: start=1 end={{ubound}}
       register: first_run
 
     - name: Reuse them
@@ -20,14 +24,16 @@
       become: true
       vars:
         ansible_become_user: "mitogen__user{{item}}"
-      with_sequence: start=1 end=21
+      with_sequence: start=1 end={{ubound}}
       register: second_run
 
     - assert:
         that:
         - first_run.results[item|int].pid == second_run.results[item|int].pid
-      with_items: start=0 end=20
+      with_items: start=0 end={{max_interps}}
+      when: is_mitogen
 
     - assert:
         that:
         - first_run.results[-1].pid != second_run.results[-1].pid
+      when: is_mitogen
diff --git a/tests/ansible/run_ansible_playbook.sh b/tests/ansible/run_ansible_playbook.sh
index 36af1755..45a4a112 100755
--- a/tests/ansible/run_ansible_playbook.sh
+++ b/tests/ansible/run_ansible_playbook.sh
@@ -3,6 +3,7 @@
 
 # Used by delegate_to.yml to ensure "sudo -E" preserves environment.
 export I_WAS_PRESERVED=1
+export MITOGEN_MAX_INTERPRETERS=3
 
 if [ "${ANSIBLE_STRATEGY:0:7}" = "mitogen" ]
 then