Merge PR239: set RLIMIT_NOFILE higher on server

closes #238

src/wormhole/server/server.py

@@ -2,10 +2,16 @@
 # a str on Python 2
 from __future__ import print_function
 import os, time, json
+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
 from twisted.internet import reactor, endpoints
 from twisted.application import service, internet
-from twisted.web import server, static, resource
+from twisted.web import server, static
+from twisted.web.resource import Resource
 from autobahn.twisted.resource import WebSocketResource
 from .database import get_db
 from .rendezvous import Rendezvous
@@ -18,10 +24,10 @@ MINUTE = 60*SECONDS
 CHANNEL_EXPIRATION_TIME = 11*MINUTE
 EXPIRATION_CHECK_PERIOD = 10*MINUTE
 
-class Root(resource.Resource):
+class Root(Resource):
     # child_FOO is a nevow thing, not a twisted.web.resource thing
     def __init__(self):
-        resource.Resource.__init__(self)
+        Resource.__init__(self)
         self.putChild(b"", static.Data(b"Wormhole Relay\n", "text/plain"))
 
 class PrivacyEnhancedSite(server.Site):
@@ -102,8 +108,38 @@ class RelayServer(service.MultiService):
             self._transit = transit
             self._transit_service = transit_service
 
+    def increase_rlimits(self):
+        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")
+
     def startService(self):
         service.MultiService.startService(self)
+        self.increase_rlimits()
         log.msg("websocket listening on /wormhole-relay/ws")
         log.msg("Wormhole relay server (Rendezvous and Transit) running")
         if self._blur_usage:

src/wormhole/test/test_server.py

@@ -3,7 +3,7 @@ import os, json, itertools, time
 import mock
 from twisted.trial import unittest
 from twisted.python import log
-from twisted.internet import reactor, defer
+from twisted.internet import reactor, defer, endpoints
 from twisted.internet.defer import inlineCallbacks, returnValue
 from autobahn.twisted import websocket
 from .common import ServerBase
@@ -24,6 +24,65 @@ def easy_relay(
         **kwargs
     )
 
+class RLimits(unittest.TestCase):
+    def test_rlimit(self):
+        def patch_s(name, *args, **kwargs):
+            return mock.patch("wormhole.server.server." + name, *args, **kwargs)
+        # We never start this, so the endpoints can be fake.
+        # serverFromString() requires bytes on py2 and str on py3, so this
+        # is easier than just passing "tcp:0"
+        ep = endpoints.TCP4ServerEndpoint(None, 0)
+        with patch_s("endpoints.serverFromString", return_value=ep):
+            s = server.RelayServer("fake", None, None)
+        fakelog = []
+        def checklog(*expected):
+            self.assertEqual(fakelog, list(expected))
+            fakelog[:] = []
+        NF = "NOFILE"
+        mock_NF = patch_s("RLIMIT_NOFILE", NF)
+
+        with patch_s("log.msg", fakelog.append):
+            with patch_s("getrlimit", None):
+                s.increase_rlimits()
+            checklog("unable to import 'resource', leaving rlimit alone")
+
+            with mock_NF:
+                with patch_s("getrlimit", return_value=(20000, 30000)) as gr:
+                    s.increase_rlimits()
+                    self.assertEqual(gr.mock_calls, [mock.call(NF)])
+                    checklog("RLIMIT_NOFILE.soft was 20000, leaving it alone")
+
+                with patch_s("getrlimit", return_value=(10, 30000)) as gr:
+                    with patch_s("setrlimit", side_effect=TypeError("other")):
+                        with patch_s("log.err") as err:
+                            s.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_s("setrlimit", side_effect=setrlimit) as sr:
+                            s.increase_rlimits()
+                        self.assertEqual(sr.mock_calls, calls)
+                        checklog(*expected)
+
 class _Util:
     def _nameplate(self, app, name):
         np_row = app._db.execute("SELECT * FROM `nameplates`"