prolog.engine: use a urllib3 HTTP connection pool

Opening large numbers of single-shot requests caused local port
exhaustion due to TIME_WAIT.

1 files changed, 35 insertions, 35 deletions

diff --git a/prolog/engine.py b/prolog/engine.py
index 780799e..e0f8bf7 100644
--- a/prolog/engine.py
+++ b/prolog/engine.py
@@ -22,7 +22,7 @@ import json
 from operator import itemgetter
 import re
 import time
-import urllib
+import urllib3
 
 def strip_html(text):
     return html.unescape(re.sub(r'</?[a-z]+[^>]*>', '', text))
@@ -30,8 +30,8 @@ def strip_html(text):
 # Create a new pengine and initialize it with [code]. Return engine ID and a
 # list of messages from Prolog.
 def create(code='', timeout=10):
-    opts = {'format': 'json-html', 'destroy': False, 'src_text': code}
-    reply, output = request('POST', '/pengine/create', body=json.dumps(opts), timeout=timeout)
+    opts = json.dumps({'format': 'json-html', 'destroy': False, 'src_text': code})
+    reply, output = request('POST', '/pengine/create', params=opts, timeout=timeout)
     return reply.get('id'), output
 
 def ask(engine, query, timeout=10):
@@ -46,47 +46,39 @@ def stop(engine, timeout=10):
     return send(engine, 'stop', timeout=timeout)
 
 def destroy(engine):
-    params = urllib.parse.urlencode({'ids': engine})
     try:
         # We don't care about the answer here, so don't wait for it.
-        request('GET', '/pengine/destroy_all?' + params, timeout=0.01)
+        params = {'ids': engine}
+        request('GET', '/pengine/destroy_all', params=params, timeout=0.01)
     except:
         pass
 
 def send(engine, event, timeout=10):
-    params = urllib.parse.urlencode({
-        'id': engine,
-        'event': event,
-        'format': 'json-html'})
-    return request('GET', path='/pengine/send?' + params, timeout=timeout)
+    params = {'id': engine, 'event': event, 'format': 'json-html'}
+    return request('GET', path='/pengine/send', params=params, timeout=timeout)
 
 # Return the main reply and pull potential output replies.
-address, port = '127.0.0.1', 3030  # TODO put this somewhere sane
-def request(method, path, body=None, timeout=10):
-    headers = {'Content-Type': 'application/json;charset=utf-8'}
+def request(method, path, params=None, timeout=10):
+    if method == 'GET':
+        response = conn.request_encode_url(method, path, fields=params, timeout=timeout)
+    else:
+        response = conn.urlopen(method, path, body=params, timeout=timeout)
+
     messages = []
-    try:
-        conn = http.client.HTTPConnection(address, port, timeout=timeout)
-        conn.request(method, path, body, headers=headers)
-        while True:
-            response = conn.getresponse()
-            if response.status != http.client.OK:
-                raise Exception('server returned {}'.format(response.status))
-
-            reply = json.loads(response.read().decode('utf-8'))
-            if isinstance(reply, dict) and reply.get('event') == 'output':
-                messages.append(_get_message(reply))
-
-                # Pull the next output. These requests should return instantly
-                # as no additional processing needs to be done in the pengine.
-                params = urllib.parse.urlencode({
-                    'id': reply['id'],
-                    'format': 'json-html'})
-                conn.request('GET', '/pengine/pull_response?' + params, headers=headers)
-            else:
-                return reply, messages
-    finally:
-        conn.close()
+    while True:
+        if response.status != http.client.OK:
+            raise Exception('server returned {}'.format(response.status))
+
+        reply = json.loads(response.data.decode('utf-8'))
+        if isinstance(reply, dict) and reply.get('event') == 'output':
+            messages.append(_get_message(reply))
+
+            # Pull the next output. These requests should return instantly
+            # as no additional processing needs to be done in the pengine.
+            params = {'id': reply['id'], 'format': 'json-html'}
+            response = conn.request_encode_url('GET', '/pengine/pull_response', fields=params)
+        else:
+            return reply, messages
 
 # Strip boilerplate from Prolog messages … ugly.
 def _get_message(reply):
@@ -240,6 +232,14 @@ def process_answer(reply):
             error = ('error', reply['data'].replace("'{}':".format(reply['id']), ''))
             return None, error, False
 
+address, port = '127.0.0.1', 3030  # TODO put this somewhere sane
+conn = urllib3.connectionpool.HTTPConnectionPool(
+        host=address, port=port,
+        timeout=10,
+        maxsize=100,
+        headers={'Content-Type': 'application/json;charset=utf-8'},
+        retries=False)
+
 # Basic sanity check.
 if __name__ == '__main__':
     engine, messages = create('b(Y). a(X) :- {X > 3}, (X = 5 ; {X > 4}).')