d0a40ee708ad074a9cce9049c966aa51f312f4d3
[quix0rs-apt-p2p.git] / apt_p2p_Khashmir / DHT.py
1
2 """The main interface to the Khashmir DHT.
3
4 @var khashmir_dir: the name of the directory to use for DHT files
5 """
6
7 from datetime import datetime
8 from StringIO import StringIO
9 import os, sha, random
10
11 from twisted.internet import defer, reactor
12 from twisted.internet.abstract import isIPAddress
13 from twisted.python import log
14 from twisted.trial import unittest
15 from zope.interface import implements
16
17 from apt_p2p.interfaces import IDHT, IDHTStats, IDHTStatsFactory
18 from khashmir import Khashmir
19 from bencode import bencode, bdecode
20
21 try:
22     from twisted.web2 import channel, server, resource, http, http_headers
23     _web2 = True
24 except ImportError:
25     _web2 = False
26
27 khashmir_dir = 'apt-p2p-Khashmir'
28
29 class DHTError(Exception):
30     """Represents errors that occur in the DHT."""
31
32 class DHT:
33     """The main interface instance to the Khashmir DHT.
34     
35     @type config: C{dictionary}
36     @ivar config: the DHT configuration values
37     @type cache_dir: C{string}
38     @ivar cache_dir: the directory to use for storing files
39     @type bootstrap: C{list} of C{string}
40     @ivar bootstrap: the nodes to contact to bootstrap into the system
41     @type bootstrap_node: C{boolean}
42     @ivar bootstrap_node: whether this node is a bootstrap node
43     @type joining: L{twisted.internet.defer.Deferred}
44     @ivar joining: if a join is underway, the deferred that will signal it's end
45     @type joined: C{boolean}
46     @ivar joined: whether the DHT network has been successfully joined
47     @type outstandingJoins: C{int}
48     @ivar outstandingJoins: the number of bootstrap nodes that have yet to respond
49     @type foundAddrs: C{list} of (C{string}, C{int})
50     @ivar foundAddrs: the IP address an port that were returned by bootstrap nodes
51     @type storing: C{dictionary}
52     @ivar storing: keys are keys for which store requests are active, values
53         are dictionaries with keys the values being stored and values the
54         deferred to call when complete
55     @type retrieving: C{dictionary}
56     @ivar retrieving: keys are the keys for which getValue requests are active,
57         values are lists of the deferreds waiting for the requests
58     @type retrieved: C{dictionary}
59     @ivar retrieved: keys are the keys for which getValue requests are active,
60         values are list of the values returned so far
61     @type factory: L{twisted.web2.channel.HTTPFactory}
62     @ivar factory: the factory to use to serve HTTP requests for statistics
63     @type config_parser: L{apt_p2p.apt_p2p_conf.AptP2PConfigParser}
64     @ivar config_parser: the configuration info for the main program
65     @type section: C{string}
66     @ivar section: the section of the configuration info that applies to the DHT
67     @type khashmir: L{khashmir.Khashmir}
68     @ivar khashmir: the khashmir DHT instance to use
69     """
70     
71     if _web2:
72         implements(IDHT, IDHTStats, IDHTStatsFactory)
73     else:
74         implements(IDHT, IDHTStats)
75         
76     def __init__(self):
77         """Initialize the DHT."""
78         self.config = None
79         self.cache_dir = ''
80         self.bootstrap = []
81         self.bootstrap_node = False
82         self.joining = None
83         self.joined = False
84         self.outstandingJoins = 0
85         self.foundAddrs = []
86         self.storing = {}
87         self.retrieving = {}
88         self.retrieved = {}
89         self.factory = None
90     
91     def loadConfig(self, config, section):
92         """See L{apt_p2p.interfaces.IDHT}."""
93         self.config_parser = config
94         self.section = section
95         self.config = {}
96         
97         # Get some initial values
98         self.cache_dir = os.path.join(self.config_parser.get(section, 'cache_dir'), khashmir_dir)
99         if not os.path.exists(self.cache_dir):
100             os.makedirs(self.cache_dir)
101         self.bootstrap = self.config_parser.getstringlist(section, 'BOOTSTRAP')
102         self.bootstrap_node = self.config_parser.getboolean(section, 'BOOTSTRAP_NODE')
103         for k in self.config_parser.options(section):
104             # The numbers in the config file
105             if k in ['K', 'HASH_LENGTH', 'CONCURRENT_REQS', 'STORE_REDUNDANCY', 
106                      'RETRIEVE_VALUES', 'MAX_FAILURES', 'PORT']:
107                 self.config[k] = self.config_parser.getint(section, k)
108             # The times in the config file
109             elif k in ['CHECKPOINT_INTERVAL', 'MIN_PING_INTERVAL', 
110                        'BUCKET_STALENESS', 'KEY_EXPIRE']:
111                 self.config[k] = self.config_parser.gettime(section, k)
112             # The booleans in the config file
113             elif k in ['SPEW']:
114                 self.config[k] = self.config_parser.getboolean(section, k)
115             # Everything else is a string
116             else:
117                 self.config[k] = self.config_parser.get(section, k)
118     
119     def join(self):
120         """See L{apt_p2p.interfaces.IDHT}."""
121         if self.config is None:
122             raise DHTError, "configuration not loaded"
123         if self.joining:
124             raise DHTError, "a join is already in progress"
125
126         # Create the new khashmir instance
127         self.khashmir = Khashmir(self.config, self.cache_dir)
128         
129         self.joining = defer.Deferred()
130         for node in self.bootstrap:
131             host, port = node.rsplit(':', 1)
132             port = int(port)
133             
134             # Translate host names into IP addresses
135             if isIPAddress(host):
136                 self._join_gotIP(host, port)
137             else:
138                 reactor.resolve(host).addCallback(self._join_gotIP, port)
139         
140         return self.joining
141
142     def _join_gotIP(self, ip, port):
143         """Join the DHT using a single bootstrap nodes IP address."""
144         self.outstandingJoins += 1
145         self.khashmir.addContact(ip, port, self._join_single, self._join_error)
146     
147     def _join_single(self, addr):
148         """Process the response from the bootstrap node.
149         
150         Finish the join by contacting close nodes.
151         """
152         self.outstandingJoins -= 1
153         if addr:
154             self.foundAddrs.append(addr)
155         if addr or self.outstandingJoins <= 0:
156             self.khashmir.findCloseNodes(self._join_complete, self._join_complete)
157         log.msg('Got back from bootstrap node: %r' % (addr,))
158     
159     def _join_error(self, failure = None):
160         """Process an error in contacting the bootstrap node.
161         
162         If no bootstrap nodes remain, finish the process by contacting
163         close nodes.
164         """
165         self.outstandingJoins -= 1
166         log.msg("bootstrap node could not be reached")
167         if self.outstandingJoins <= 0:
168             self.khashmir.findCloseNodes(self._join_complete, self._join_complete)
169
170     def _join_complete(self, result):
171         """End the joining process and return the addresses found for this node."""
172         if not self.joined and len(result) > 0:
173             self.joined = True
174         if self.joining and self.outstandingJoins <= 0:
175             df = self.joining
176             self.joining = None
177             if self.joined or self.bootstrap_node:
178                 self.joined = True
179                 df.callback(self.foundAddrs)
180             else:
181                 df.errback(DHTError('could not find any nodes to bootstrap to'))
182         
183     def getAddrs(self):
184         """Get the list of addresses returned by bootstrap nodes for this node."""
185         return self.foundAddrs
186         
187     def leave(self):
188         """See L{apt_p2p.interfaces.IDHT}."""
189         if self.config is None:
190             raise DHTError, "configuration not loaded"
191         
192         if self.joined or self.joining:
193             if self.joining:
194                 self.joining.errback(DHTError('still joining when leave was called'))
195                 self.joining = None
196             self.joined = False
197             self.khashmir.shutdown()
198         
199     def _normKey(self, key, bits=None, bytes=None):
200         """Normalize the length of keys used in the DHT."""
201         bits = self.config["HASH_LENGTH"]
202         if bits is not None:
203             bytes = (bits - 1) // 8 + 1
204         else:
205             if bytes is None:
206                 raise DHTError, "you must specify one of bits or bytes for normalization"
207             
208         # Extend short keys with null bytes
209         if len(key) < bytes:
210             key = key + '\000'*(bytes - len(key))
211         # Truncate long keys
212         elif len(key) > bytes:
213             key = key[:bytes]
214         return key
215
216     def getValue(self, key):
217         """See L{apt_p2p.interfaces.IDHT}."""
218         if self.config is None:
219             raise DHTError, "configuration not loaded"
220         if not self.joined:
221             raise DHTError, "have not joined a network yet"
222         
223         key = self._normKey(key)
224
225         d = defer.Deferred()
226         if key not in self.retrieving:
227             self.khashmir.valueForKey(key, self._getValue)
228         self.retrieving.setdefault(key, []).append(d)
229         return d
230         
231     def _getValue(self, key, result):
232         """Process a returned list of values from the DHT."""
233         # Save the list of values to return when it is complete
234         if result:
235             self.retrieved.setdefault(key, []).extend([bdecode(r) for r in result])
236         else:
237             # Empty list, the get is complete, return the result
238             final_result = []
239             if key in self.retrieved:
240                 final_result = self.retrieved[key]
241                 del self.retrieved[key]
242             for i in range(len(self.retrieving[key])):
243                 d = self.retrieving[key].pop(0)
244                 d.callback(final_result)
245             del self.retrieving[key]
246
247     def storeValue(self, key, value):
248         """See L{apt_p2p.interfaces.IDHT}."""
249         if self.config is None:
250             raise DHTError, "configuration not loaded"
251         if not self.joined:
252             raise DHTError, "have not joined a network yet"
253
254         key = self._normKey(key)
255         bvalue = bencode(value)
256
257         if key in self.storing and bvalue in self.storing[key]:
258             raise DHTError, "already storing that key with the same value"
259
260         d = defer.Deferred()
261         self.khashmir.storeValueForKey(key, bvalue, self._storeValue)
262         self.storing.setdefault(key, {})[bvalue] = d
263         return d
264     
265     def _storeValue(self, key, bvalue, result):
266         """Process the response from the DHT."""
267         if key in self.storing and bvalue in self.storing[key]:
268             # Check if the store succeeded
269             if len(result) > 0:
270                 self.storing[key][bvalue].callback(result)
271             else:
272                 self.storing[key][bvalue].errback(DHTError('could not store value %s in key %s' % (bvalue, key)))
273             del self.storing[key][bvalue]
274             if len(self.storing[key].keys()) == 0:
275                 del self.storing[key]
276     
277     def getStats(self):
278         """See L{apt_p2p.interfaces.IDHTStats}."""
279         stats = self.khashmir.getStats()
280         out = StringIO()
281         out.write('<h2>DHT Statistics</h2>\n')
282         old_group = None
283         for stat in stats:
284             if stat['group'] != old_group:
285                 if old_group is not None:
286                     out.write('</table>\n')
287                 out.write('\n<h3>' + stat['group'] + '</h3>\n')
288                 out.write("<table border='1'>\n")
289                 if stat['group'] != 'Actions':
290                     out.write("<tr><th>Statistic</th><th>Value</th></tr>\n")
291                 else:
292                     out.write("<tr><th>Action</th><th>Sent</th><th>OK</th><th>Failed</th><th>Received</th><th>Error</th></tr>\n")
293                 old_group = stat['group']
294             if stat['group'] != 'Actions':
295                 out.write("<tr title='" + stat['tip'] + "'><td>" + stat['desc'] + '</td><td>' + str(stat['value']) + '</td></tr>\n')
296             else:
297                 actions = stat['value'].keys()
298                 actions.sort()
299                 for action in actions:
300                     out.write("<tr><td>" + action + "</td>")
301                     for i in xrange(5):
302                         out.write("<td>" + str(stat['value'][action][i]) + "</td>")
303                     out.write('</tr>\n')
304                     
305         return out.getvalue()
306
307     def getStatsFactory(self):
308         """See L{apt_p2p.interfaces.IDHTStatsFactory}."""
309         assert _web2, "NOT IMPLEMENTED: twisted.web2 must be installed to use the stats factory."
310         if self.factory is None:
311             # Create a simple HTTP factory for stats
312             class StatsResource(resource.Resource):
313                 def __init__(self, manager):
314                     self.manager = manager
315                 def render(self, ctx):
316                     return http.Response(
317                         200,
318                         {'content-type': http_headers.MimeType('text', 'html')},
319                         '<html><body>\n\n' + self.manager.getStats() + '\n</body></html>\n')
320                 def locateChild(self, request, segments):
321                     log.msg('Got HTTP stats request from %s' % (request.remoteAddr, ))
322                     return self, ()
323             
324             self.factory = channel.HTTPFactory(server.Site(StatsResource(self)))
325         return self.factory
326         
327
328 class TestSimpleDHT(unittest.TestCase):
329     """Simple 2-node unit tests for the DHT."""
330     
331     timeout = 2
332     DHT_DEFAULTS = {'PORT': 9977, 'K': 8, 'HASH_LENGTH': 160,
333                     'CHECKPOINT_INTERVAL': 300, 'CONCURRENT_REQS': 4,
334                     'STORE_REDUNDANCY': 3, 'RETRIEVE_VALUES': -10000,
335                     'MAX_FAILURES': 3,
336                     'MIN_PING_INTERVAL': 900,'BUCKET_STALENESS': 3600,
337                     'KEY_EXPIRE': 3600, 'SPEW': False, }
338
339     def setUp(self):
340         self.a = DHT()
341         self.b = DHT()
342         self.a.config = self.DHT_DEFAULTS.copy()
343         self.a.config['PORT'] = 4044
344         self.a.bootstrap = ["127.0.0.1:4044"]
345         self.a.bootstrap_node = True
346         self.a.cache_dir = '/tmp'
347         self.b.config = self.DHT_DEFAULTS.copy()
348         self.b.config['PORT'] = 4045
349         self.b.bootstrap = ["127.0.0.1:4044"]
350         self.b.cache_dir = '/tmp'
351         
352     def test_bootstrap_join(self):
353         d = self.a.join()
354         return d
355         
356     def node_join(self, result):
357         d = self.b.join()
358         return d
359     
360     def test_join(self):
361         self.lastDefer = defer.Deferred()
362         d = self.a.join()
363         d.addCallback(self.node_join)
364         d.addCallback(self.lastDefer.callback)
365         return self.lastDefer
366
367     def test_normKey(self):
368         h = self.a._normKey('12345678901234567890')
369         self.failUnless(h == '12345678901234567890')
370         h = self.a._normKey('12345678901234567')
371         self.failUnless(h == '12345678901234567\000\000\000')
372         h = self.a._normKey('1234567890123456789012345')
373         self.failUnless(h == '12345678901234567890')
374         h = self.a._normKey('1234567890123456789')
375         self.failUnless(h == '1234567890123456789\000')
376         h = self.a._normKey('123456789012345678901')
377         self.failUnless(h == '12345678901234567890')
378
379     def value_stored(self, result, value):
380         self.stored -= 1
381         if self.stored == 0:
382             self.get_values()
383         
384     def store_values(self, result):
385         self.stored = 3
386         d = self.a.storeValue(sha.new('4045').digest(), str(4045*3))
387         d.addCallback(self.value_stored, 4045)
388         d = self.a.storeValue(sha.new('4044').digest(), str(4044*2))
389         d.addCallback(self.value_stored, 4044)
390         d = self.b.storeValue(sha.new('4045').digest(), str(4045*2))
391         d.addCallback(self.value_stored, 4045)
392
393     def check_values(self, result, values):
394         self.checked -= 1
395         self.failUnless(len(result) == len(values))
396         for v in result:
397             self.failUnless(v in values)
398         if self.checked == 0:
399             self.lastDefer.callback(1)
400     
401     def get_values(self):
402         self.checked = 4
403         d = self.a.getValue(sha.new('4044').digest())
404         d.addCallback(self.check_values, [str(4044*2)])
405         d = self.b.getValue(sha.new('4044').digest())
406         d.addCallback(self.check_values, [str(4044*2)])
407         d = self.a.getValue(sha.new('4045').digest())
408         d.addCallback(self.check_values, [str(4045*2), str(4045*3)])
409         d = self.b.getValue(sha.new('4045').digest())
410         d.addCallback(self.check_values, [str(4045*2), str(4045*3)])
411
412     def test_store(self):
413         from twisted.internet.base import DelayedCall
414         DelayedCall.debug = True
415         self.lastDefer = defer.Deferred()
416         d = self.a.join()
417         d.addCallback(self.node_join)
418         d.addCallback(self.store_values)
419         return self.lastDefer
420
421     def tearDown(self):
422         self.a.leave()
423         try:
424             os.unlink(self.a.khashmir.store.db)
425         except:
426             pass
427         self.b.leave()
428         try:
429             os.unlink(self.b.khashmir.store.db)
430         except:
431             pass
432
433 class TestMultiDHT(unittest.TestCase):
434     """More complicated 20-node tests for the DHT."""
435     
436     timeout = 60
437     num = 20
438     DHT_DEFAULTS = {'PORT': 9977, 'K': 8, 'HASH_LENGTH': 160,
439                     'CHECKPOINT_INTERVAL': 300, 'CONCURRENT_REQS': 4,
440                     'STORE_REDUNDANCY': 3, 'RETRIEVE_VALUES': -10000,
441                     'MAX_FAILURES': 3,
442                     'MIN_PING_INTERVAL': 900,'BUCKET_STALENESS': 3600,
443                     'KEY_EXPIRE': 3600, 'SPEW': False, }
444
445     def setUp(self):
446         self.l = []
447         self.startport = 4081
448         for i in range(self.num):
449             self.l.append(DHT())
450             self.l[i].config = self.DHT_DEFAULTS.copy()
451             self.l[i].config['PORT'] = self.startport + i
452             self.l[i].bootstrap = ["127.0.0.1:" + str(self.startport)]
453             self.l[i].cache_dir = '/tmp'
454         self.l[0].bootstrap_node = True
455         
456     def node_join(self, result, next_node):
457         d = self.l[next_node].join()
458         if next_node + 1 < len(self.l):
459             d.addCallback(self.node_join, next_node + 1)
460         else:
461             d.addCallback(self.lastDefer.callback)
462     
463     def test_join(self):
464         self.timeout = 2
465         self.lastDefer = defer.Deferred()
466         d = self.l[0].join()
467         d.addCallback(self.node_join, 1)
468         return self.lastDefer
469         
470     def store_values(self, result, i = 0, j = 0):
471         if j > i:
472             j -= i+1
473             i += 1
474         if i == len(self.l):
475             self.get_values()
476         else:
477             d = self.l[j].storeValue(sha.new(str(self.startport+i)).digest(), str((self.startport+i)*(j+1)))
478             d.addCallback(self.store_values, i, j+1)
479     
480     def get_values(self, result = None, check = None, i = 0, j = 0):
481         if result is not None:
482             self.failUnless(len(result) == len(check))
483             for v in result:
484                 self.failUnless(v in check)
485         if j >= len(self.l):
486             j -= len(self.l)
487             i += 1
488         if i == len(self.l):
489             self.lastDefer.callback(1)
490         else:
491             d = self.l[i].getValue(sha.new(str(self.startport+j)).digest())
492             check = []
493             for k in range(self.startport+j, (self.startport+j)*(j+1)+1, self.startport+j):
494                 check.append(str(k))
495             d.addCallback(self.get_values, check, i, j + random.randrange(1, min(len(self.l), 10)))
496
497     def store_join(self, result, next_node):
498         d = self.l[next_node].join()
499         if next_node + 1 < len(self.l):
500             d.addCallback(self.store_join, next_node + 1)
501         else:
502             d.addCallback(self.store_values)
503     
504     def test_store(self):
505         from twisted.internet.base import DelayedCall
506         DelayedCall.debug = True
507         self.lastDefer = defer.Deferred()
508         d = self.l[0].join()
509         d.addCallback(self.store_join, 1)
510         return self.lastDefer
511
512     def tearDown(self):
513         for i in self.l:
514             try:
515                 i.leave()
516                 os.unlink(i.khashmir.store.db)
517             except:
518                 pass