return K nodes and not K-1

[quix0rs-apt-p2p.git] / ktable.py

diff --git a/ktable.py b/ktable.py
index 3bff4074ff75a1fcf5be33fe3f6d6bad6ef67e3b..f1c3b6ac50fa5856d2cdc96dd8ab725c0756df8d 100644 (file)
--- a/ktable.py
+++ b/ktable.py
@@ -5,6 +5,7 @@ from bisect import *
 import time
 from types import *
 
+import const
 from node import Node
 
 # The all-powerful, magical Kademlia "k" constant, bucket depth
@@ -64,23 +65,24 @@ class KTable:
            return [self.buckets[i].l[index]]
            
        nodes = nodes + self.buckets[i].l
-       if len(nodes) == K:
+       if len(nodes) >= K:
            nodes.sort(sort)
-           return nodes
+           return nodes[:K]
        else:
            # need more nodes
            min = i - 1
            max = i + 1
-           while (len(nodes) < K and (min >= 0 and max < len(self.buckets))):
+           while (len(nodes) < K and (min >= 0 or max < len(self.buckets))):
                if min >= 0:
                    nodes = nodes + self.buckets[min].l
                    self.buckets[min].touch()
                if max < len(self.buckets):
                    nodes = nodes + self.buckets[max].l
                    self.buckets[max].touch()
-               
+               min = min - 1
+               max = max + 1
            nodes.sort(sort)
-           return nodes[:K-1]
+           return nodes[:K]
 
     def _splitBucket(self, a):
        diff = (a.max - a.min) / 2
@@ -104,12 +106,14 @@ class KTable:
            return
 
        del(self.buckets[i].l[it])
-       self.buckets[i].l.append(new)
+       if new:
+           self.buckets[i].l.append(new)
 
-    def insertNode(self, node):
+    def insertNode(self, node, contacted=1):
        """ 
        this insert the node, returning None if successful, returns the oldest node in the bucket if it's full
        the caller responsible for pinging the returned node and calling replaceStaleNode if it is found to be stale!!
+       contacted means that yes, we contacted THEM and we know the node is available
        """
        assert(node.id != " "*20)
        # get the bucket for this node
@@ -121,16 +125,21 @@ class KTable:
            ## no
            pass
        else:
-           node.updateLastSeen()
-           # move node to end of bucket
-           del(self.buckets[i].l[it])
-           self.buckets[i].l.append(node)
-           self.buckets[i].touch()
+           if contacted:
+               node.updateLastSeen()
+               # move node to end of bucket
+               del(self.buckets[i].l[it])
+               # note that we removed the original and replaced it with the new one
+               # utilizing this nodes new contact info
+               self.buckets[i].l.append(node)
+               self.buckets[i].touch()
            return
        
        # we don't have this node, check to see if the bucket is full
        if len(self.buckets[i].l) < K:
            # no, append this node and return
+           if contacted:
+               node.updateLastSeen()
            self.buckets[i].l.append(node)
            self.buckets[i].touch()
            return
@@ -163,7 +172,16 @@ class KTable:
            n.updateLastSeen()
            return tstamp
 
-
+    def nodeFailed(self, node):
+       """ call this when a node fails to respond to a message, to invalidate that node """
+       try:
+           n = self.findNodes(node.int)[0]
+       except IndexError:
+           return None
+       else:
+           if(n.msgFailed() >= const.MAX_FAILURES):
+               self.replaceStaleNode(n, None)
+       
 class KBucket:
     __slots = ['min', 'max', 'lastAccessed']
     def __init__(self, contents, min, max):
@@ -219,11 +237,11 @@ import unittest
 
 class TestKTable(unittest.TestCase):
     def setUp(self):
-       self.a = Node(hash.newID(), 'localhost', 2002)
+       self.a = Node().init(hash.newID(), 'localhost', 2002)
        self.t = KTable(self.a)
 
     def test_replace_stale_node(self):
-       self.b = Node(hash.newID(), 'localhost', 2003)
+       self.b = Node().init(hash.newID(), 'localhost', 2003)
        self.t.replaceStaleNode(self.a, self.b)
        assert(len(self.t.buckets[0].l) == 1)
        assert(self.t.buckets[0].l[0].id == self.b.id)