key_generator: service related cleanups, incorporation into system test

this cleans up KeyGenerator to be a service (a subservice of the
KeyGeneratorService as instantiated by the key-generator.tac app)
this means that the timer which replenishes the keypool will be
shutdown cleanly when the service is stopped.

adds checks on the key_generator service and client into the system
test 'test_mutable' such that one of the nodes (clients[3]) uses
the key_generator service, and checks that mutable file creation
in that node, via a variety of means, are all consuming keys from
the key_generator.

src/allmydata/key_generator.py

@@ -3,7 +3,6 @@ import os
 import time
 
 import foolscap
-from foolscap.eventual import eventually
 from zope.interface import implements
 from twisted.internet import reactor
 from twisted.application import service
@@ -12,7 +11,7 @@ from twisted.python import log
 from pycryptopp.publickey import rsa
 from allmydata.interfaces import RIKeyGenerator
 
-class KeyGenerator(foolscap.Referenceable):
+class KeyGenerator(service.MultiService, foolscap.Referenceable):
     implements(RIKeyGenerator)
 
     DEFAULT_KEY_SIZE = 2048
@@ -21,9 +20,18 @@ class KeyGenerator(foolscap.Referenceable):
     verbose = False
 
     def __init__(self):
+        service.MultiService.__init__(self)
         self.keypool = []
         self.last_fetch = 0
-        eventually(self.maybe_refill_pool)
+
+    def startService(self):
+        self.timer = reactor.callLater(0, self.maybe_refill_pool)
+        return service.MultiService.startService(self)
+
+    def stopService(self):
+        if self.timer.active():
+            self.timer.cancel()
+        return service.MultiService.stopService(self)
 
     def __repr__(self):
         return '<KeyGenerator[%s]>' % (len(self.keypool),)
@@ -34,7 +42,10 @@ class KeyGenerator(foolscap.Referenceable):
 
     def reset_timer(self):
         self.last_fetch = time.time()
-        reactor.callLater(self.pool_refresh_delay, self.maybe_refill_pool)
+        if self.timer.active():
+            self.timer.reset(self.pool_refresh_delay)
+        else:
+            self.timer = reactor.callLater(self.pool_refresh_delay, self.maybe_refill_pool)
 
     def maybe_refill_pool(self):
         now = time.time()
@@ -63,11 +74,13 @@ class KeyGenerator(foolscap.Referenceable):
 class KeyGeneratorService(service.MultiService):
     furl_file = 'key_generator.furl'
 
-    def __init__(self, display_furl=True):
+    def __init__(self, basedir='.', display_furl=True):
         service.MultiService.__init__(self)
-        self.tub = foolscap.Tub(certFile='key_generator.pem')
+        self.basedir = basedir
+        self.tub = foolscap.Tub(certFile=os.path.join(self.basedir, 'key_generator.pem'))
         self.tub.setServiceParent(self)
         self.key_generator = KeyGenerator()
+        self.key_generator.setServiceParent(self)
 
         portnum = self.get_portnum()
         self.listener = self.tub.listenOn(portnum or 'tcp:0')
@@ -78,14 +91,17 @@ class KeyGeneratorService(service.MultiService):
         d.addErrback(log.err)
 
     def get_portnum(self):
-        if os.path.exists('portnum'):
-            return file('portnum', 'rb').read().strip()
+        portnumfile = os.path.join(self.basedir, 'portnum')
+        if os.path.exists(portnumfile):
+            return file(portnumfile, 'rb').read().strip()
 
     def save_portnum(self, junk):
         portnum = self.listener.getPortnum()
-        file('portnum', 'wb').write('%d\n' % (portnum,))
+        portnumfile = os.path.join(self.basedir, 'portnum')
+        file(portnumfile, 'wb').write('%d\n' % (portnum,))
 
     def tub_ready(self, junk, display_furl):
-        self.keygen_furl = self.tub.registerReference(self.key_generator, furlFile=self.furl_file)
+        kgf = os.path.join(self.basedir, self.furl_file)
+        self.keygen_furl = self.tub.registerReference(self.key_generator, furlFile=kgf)
         if display_furl:
-            print 'key generator at:', self.keygen_furl 
+            print 'key generator at:', self.keygen_furl

src/allmydata/test/test_keygen.py

@@ -52,7 +52,7 @@ class KeyGenService(unittest.TestCase, testutil.PollMixin):
         d = eventual.fireEventually()
         d.addCallback(p, 'waiting for pool to fill up')
         d.addCallback(lambda junk: self.poll(keypool_full, timeout=16))
-        
+
         d.addCallback(p, 'grabbing a few keys')
         # grab a few keys, check that pool size shrinks
         def get_key(junk=None):
@@ -89,7 +89,7 @@ class KeyGenService(unittest.TestCase, testutil.PollMixin):
         # and check it still works (will gen key synchronously on demand)
         d.addCallback(get_key)
         d.addCallback(check_key_works)
-        
+
         d.addCallback(p, 'checking pool replenishment')
         # check that the pool will refill
         timeout = 2*kgs.key_generator.pool_size + kgs.key_generator.pool_refresh_delay

src/allmydata/test/test_system.py

@@ -16,7 +16,8 @@ from allmydata.scripts import runner
 from allmydata.interfaces import IDirectoryNode, IFileNode, IFileURI
 from allmydata.mutable import NotMutableError
 from allmydata.stats import PickleStatsGatherer
-from foolscap.eventual import flushEventualQueue
+from allmydata.key_generator import KeyGeneratorService
+from foolscap.eventual import flushEventualQueue, fireEventually
 from foolscap import DeadReferenceError, Tub
 from twisted.python.failure import Failure
 from twisted.web.client import getPage
@@ -80,6 +81,7 @@ class SystemTest(testutil.SignalMixin, testutil.PollMixin, unittest.TestCase):
         d = self.introducer.when_tub_ready()
         d.addCallback(self._get_introducer_web)
         d.addCallback(self._set_up_stats_gatherer)
+        d.addCallback(self._set_up_key_generator)
         d.addCallback(self._set_up_nodes_2)
         d.addCallback(self._grab_stats)
         return d
@@ -102,6 +104,25 @@ class SystemTest(testutil.SignalMixin, testutil.PollMixin, unittest.TestCase):
         self.add_service(self.stats_gatherer)
         self.stats_gatherer_furl = self.stats_gatherer.get_furl()
 
+    def _set_up_key_generator(self, res):
+        kgsdir = self.getdir("key_generator")
+        fileutil.make_dirs(kgsdir)
+
+        self.key_generator_svc = KeyGeneratorService(kgsdir, display_furl=False)
+        self.key_generator_svc.key_generator.pool_size = 4
+        self.key_generator_svc.key_generator.pool_refresh_delay = 60
+        self.add_service(self.key_generator_svc)
+
+        d = fireEventually()
+        def check_for_furl():
+            return os.path.exists(os.path.join(kgsdir, 'key_generator.furl'))
+        d.addCallback(lambda junk: self.poll(check_for_furl, timeout=30))
+        def get_furl(junk):
+            kgf = os.path.join(kgsdir, 'key_generator.furl')
+            self.key_generator_furl = file(kgf, 'rb').read().strip()
+        d.addCallback(get_furl)
+        return d
+
     def _set_up_nodes_2(self, res):
         q = self.introducer
         self.introducer_furl = q.introducer_url
@@ -112,12 +133,14 @@ class SystemTest(testutil.SignalMixin, testutil.PollMixin, unittest.TestCase):
             basedirs.append(basedir)
             fileutil.make_dirs(basedir)
             if i == 0:
-                # client[0] runs a webserver and a helper
+                # client[0] runs a webserver and a helper, no key_generator
                 open(os.path.join(basedir, "webport"), "w").write("tcp:0:interface=127.0.0.1")
                 open(os.path.join(basedir, "run_helper"), "w").write("yes\n")
             if i == 3:
-                # client[3] runs a webserver and uses a helper
+                # client[3] runs a webserver and uses a helper, uses key_generator
                 open(os.path.join(basedir, "webport"), "w").write("tcp:0:interface=127.0.0.1")
+                kgf = "%s\n" % (self.key_generator_furl,)
+                open(os.path.join(basedir, "key_generator.furl"), "w").write(kgf)
             if self.createprivdir:
                 fileutil.make_dirs(os.path.join(basedir, "private"))
                 open(os.path.join(basedir, "private", "root_dir.cap"), "w")
@@ -801,6 +824,23 @@ class SystemTest(testutil.SignalMixin, testutil.PollMixin, unittest.TestCase):
             return d1
         d.addCallback(_created_dirnode)
 
+        def wait_for_c3_kg_conn():
+            return self.clients[3]._key_generator is not None
+        d.addCallback(lambda junk: self.poll(wait_for_c3_kg_conn))
+
+        def check_kg_poolsize(junk, size_delta):
+            self.failUnlessEqual(len(self.key_generator_svc.key_generator.keypool),
+                                 self.key_generator_svc.key_generator.pool_size + size_delta)
+
+        d.addCallback(check_kg_poolsize, 0)
+        d.addCallback(lambda junk: self.clients[3].create_mutable_file('hello, world'))
+        d.addCallback(check_kg_poolsize, -1)
+        d.addCallback(lambda junk: self.clients[3].create_empty_dirnode())
+        d.addCallback(check_kg_poolsize, -2)
+        # use_helper induces use of clients[3], which is the using-key_gen client
+        d.addCallback(lambda junk: self.POST("uri", use_helper=True, t="mkdir", name='george'))
+        d.addCallback(check_kg_poolsize, -3)
+
         return d
     # The default 120 second timeout went off when running it under valgrind
     # on my old Windows laptop, so I'm bumping up the timeout.