From f7b8c5b19a762a19b38e9298cbf46203eb5f846e Mon Sep 17 00:00:00 2001
From: Brian Warner <warner@lothar.com>
Date: Mon, 19 Feb 2018 12:26:10 -0800
Subject: [PATCH] copy increase_rlimits.py from magic-wormhole-mailbox-server

This code was originally in magic-wormhole, embedded in the server startup
code. I carved it out to a separate file (with tests) in mailbox-server, but
it should be used here too.
---
 .../increase_rlimits.py                       | 35 ++++++++++++
 .../test/test_rlimits.py                      | 57 +++++++++++++++++++
 2 files changed, 92 insertions(+)
 create mode 100644 src/wormhole_transit_relay/increase_rlimits.py
 create mode 100644 src/wormhole_transit_relay/test/test_rlimits.py

diff --git a/src/wormhole_transit_relay/increase_rlimits.py b/src/wormhole_transit_relay/increase_rlimits.py
new file mode 100644
index 0000000..4cd5029
--- /dev/null
+++ b/src/wormhole_transit_relay/increase_rlimits.py
@@ -0,0 +1,35 @@
+try:
+    # 'resource' is unix-only
+    from resource import getrlimit, setrlimit, RLIMIT_NOFILE
+except ImportError: # pragma: nocover
+    getrlimit, setrlimit, RLIMIT_NOFILE = None, None, None # pragma: nocover
+from twisted.python import log
+
+def increase_rlimits():
+    if getrlimit is None:
+        log.msg("unable to import 'resource', leaving rlimit alone")
+        return
+    soft, hard = getrlimit(RLIMIT_NOFILE)
+    if soft >= 10000:
+        log.msg("RLIMIT_NOFILE.soft was %d, leaving it alone" % soft)
+        return
+    # OS-X defaults to soft=7168, and reports a huge number for 'hard',
+    # but won't accept anything more than soft=10240, so we can't just
+    # set soft=hard. Linux returns (1024, 1048576) and is fine with
+    # soft=hard. Cygwin is reported to return (256,-1) and accepts up to
+    # soft=3200. So we try multiple values until something works.
+    for newlimit in [hard, 10000, 3200, 1024]:
+        log.msg("changing RLIMIT_NOFILE from (%s,%s) to (%s,%s)" %
+                (soft, hard, newlimit, hard))
+        try:
+            setrlimit(RLIMIT_NOFILE, (newlimit, hard))
+            log.msg("setrlimit successful")
+            return
+        except ValueError as e:
+            log.msg("error during setrlimit: %s" % e)
+            continue
+        except:
+            log.msg("other error during setrlimit, leaving it alone")
+            log.err()
+            return
+    log.msg("unable to change rlimit, leaving it alone")
diff --git a/src/wormhole_transit_relay/test/test_rlimits.py b/src/wormhole_transit_relay/test/test_rlimits.py
new file mode 100644
index 0000000..10497e4
--- /dev/null
+++ b/src/wormhole_transit_relay/test/test_rlimits.py
@@ -0,0 +1,57 @@
+from __future__ import print_function, unicode_literals
+import mock
+from twisted.trial import unittest
+from ..increase_rlimits import increase_rlimits
+
+class RLimits(unittest.TestCase):
+    def test_rlimit(self):
+        def patch_r(name, *args, **kwargs):
+            return mock.patch("wormhole_transit_relay.increase_rlimits." + name, *args, **kwargs)
+        fakelog = []
+        def checklog(*expected):
+            self.assertEqual(fakelog, list(expected))
+            fakelog[:] = []
+        NF = "NOFILE"
+        mock_NF = patch_r("RLIMIT_NOFILE", NF)
+
+        with patch_r("log.msg", fakelog.append):
+            with patch_r("getrlimit", None):
+                increase_rlimits()
+            checklog("unable to import 'resource', leaving rlimit alone")
+
+            with mock_NF:
+                with patch_r("getrlimit", return_value=(20000, 30000)) as gr:
+                    increase_rlimits()
+                    self.assertEqual(gr.mock_calls, [mock.call(NF)])
+                    checklog("RLIMIT_NOFILE.soft was 20000, leaving it alone")
+
+                with patch_r("getrlimit", return_value=(10, 30000)) as gr:
+                    with patch_r("setrlimit", side_effect=TypeError("other")):
+                        with patch_r("log.err") as err:
+                            increase_rlimits()
+                        self.assertEqual(err.mock_calls, [mock.call()])
+                        checklog("changing RLIMIT_NOFILE from (10,30000) to (30000,30000)",
+                                 "other error during setrlimit, leaving it alone")
+
+                    for maxlimit in [40000, 20000, 9000, 2000, 1000]:
+                        def setrlimit(which, newlimit):
+                            if newlimit[0] > maxlimit:
+                                raise ValueError("nope")
+                            return None
+                        calls = []
+                        expected = []
+                        for tries in [30000, 10000, 3200, 1024]:
+                            calls.append(mock.call(NF, (tries, 30000)))
+                            expected.append("changing RLIMIT_NOFILE from (10,30000) to (%d,30000)" % tries)
+                            if tries > maxlimit:
+                                expected.append("error during setrlimit: nope")
+                            else:
+                                expected.append("setrlimit successful")
+                                break
+                        else:
+                            expected.append("unable to change rlimit, leaving it alone")
+
+                        with patch_r("setrlimit", side_effect=setrlimit) as sr:
+                            increase_rlimits()
+                        self.assertEqual(sr.mock_calls, calls)
+                        checklog(*expected)