From a6818c4300e7cb7628b79a7a6a56dfd49a1d0e5b Mon Sep 17 00:00:00 2001
From: zelig <viktor.tron@gmail.com>
Date: Mon, 11 May 2015 19:55:06 +0200
Subject: [PATCH] kademlia: added callback on short bucket + tests

---
 common/kademlia/kademlia.go      |  7 +++-
 common/kademlia/kademlia_test.go | 55 ++++++++++++++++++++++++++++----
 2 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/common/kademlia/kademlia.go b/common/kademlia/kademlia.go
index 9309bc3133..fe0d2b4b27 100644
--- a/common/kademlia/kademlia.go
+++ b/common/kademlia/kademlia.go
@@ -34,6 +34,8 @@ type Kademlia struct {
 	nodeDB         [][]*nodeRecord
 	nodeIndex      map[Address]*nodeRecord
 
+	GetNode func(int)
+
 	// state
 	proxLimit int
 	proxSize  int
@@ -52,6 +54,7 @@ type Node interface {
 	Addr() Address
 	// Url()
 	LastActive() time.Time
+	Drop()
 }
 
 type nodeRecord struct {
@@ -159,7 +162,9 @@ func (self *Kademlia) RemoveNode(node Node) (err error) {
 	}
 	// async callback to notify user that bucket needs filling
 	// action is left to the user
-	// go self.getNode(index)
+	if self.GetNode != nil {
+		go self.GetNode(index)
+	}
 	return
 }
 
diff --git a/common/kademlia/kademlia_test.go b/common/kademlia/kademlia_test.go
index 9de7a3138f..3e05c0ec51 100644
--- a/common/kademlia/kademlia_test.go
+++ b/common/kademlia/kademlia_test.go
@@ -11,7 +11,6 @@ import (
 	"testing/quick"
 	"time"
 
-	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/logger"
 )
 
@@ -40,6 +39,9 @@ func (n *testNode) Addr() Address {
 	return n.addr
 }
 
+func (n *testNode) Drop() {
+}
+
 func (n *testNode) LastActive() time.Time {
 	return time.Now()
 }
@@ -154,12 +156,10 @@ func TestProxAdjust(t *testing.T) {
 		a := gen(Address{}, r).(Address)
 		addresses = append(addresses, a)
 		err = kad.AddNode(&testNode{addr: a})
-		fmt.Printf("add node: %x (%v)\n", a, kad.proximityBin(a))
 		if err != nil {
 			t.Errorf("backend not accepting node")
 			return
 		}
-		fmt.Printf("MaxProxBinSize: %d, proxSize: %d, proxLimit: %d, count: %d\n", kad.MaxProxBinSize, kad.proxSize, kad.proxLimit, kad.count)
 		if !kad.proxCheck(t) {
 			return
 		}
@@ -167,15 +167,11 @@ func TestProxAdjust(t *testing.T) {
 
 	test := func(test *proxTest) bool {
 		node := &testNode{test.address}
-		a := test.address
 		if test.add {
 			kad.AddNode(node)
-			fmt.Printf("add node: %x (%v)\n", a, kad.proximityBin(a))
 		} else {
 			kad.RemoveNode(node)
-			fmt.Printf("remove node: %x (%v)\n", common.ToHex(a[:]), kad.proximityBin(a))
 		}
-		fmt.Printf("MaxProxBinSize: %d, proxSize: %d, proxLimit: %d, count: %d\n", kad.MaxProxBinSize, kad.proxSize, kad.proxLimit, kad.count)
 		return kad.proxCheck(t)
 	}
 	if err := quick.Check(test, quickcfg); err != nil {
@@ -183,6 +179,51 @@ func TestProxAdjust(t *testing.T) {
 	}
 }
 
+func TestCallback(t *testing.T) {
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	self := gen(Address{}, r).(Address)
+	kad := New(self)
+	var bucket int
+	var called chan bool
+	kad.MaxProx = 5
+	kad.BucketSize = 2
+	kad.GetNode = func(b int) {
+		bucket = b
+		close(called)
+	}
+	kad.Start()
+	var err error
+
+	for i := 0; i < 100; i++ {
+		a := gen(Address{}, r).(Address)
+		addresses = append(addresses, a)
+		err = kad.AddNode(&testNode{addr: a})
+		if err != nil {
+			t.Errorf("backend not accepting node")
+			return
+		}
+	}
+	for _, a := range addresses {
+		called = make(chan bool)
+		kad.RemoveNode(&testNode{a})
+		b := kad.proximityBin(a)
+		select {
+		case <-called:
+			if bucket != b {
+				t.Errorf("GetNode callback called with incorrect bucket, expected %v, got %v", b, bucket)
+			}
+		case <-time.After(100 * time.Millisecond):
+			l := len(kad.buckets[b].nodes)
+			if l < kad.BucketSize {
+				t.Errorf("GetNode not called on bucket %v although size is %v < %v", b, l, kad.BucketSize)
+			} else {
+				t.Logf("bucket callback ok")
+			}
+		}
+	}
+
+}
+
 func TestSaveLoad(t *testing.T) {
 	r := rand.New(rand.NewSource(time.Now().UnixNano()))
 	addresses := gen([]Address{}, r).([]Address)