WIP

probe-lab · Oct 5, 2023 · b55f2db · b55f2db
1 parent 5c6e8a6
commit b55f2db
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 19 deletions.
diff --git a/internal/coord/query/crawl.go → internal/coord/routing/crawl.go b/internal/coord/query/crawl.go → internal/coord/routing/crawl.go
@@ -1,4 +1,4 @@
-package query
+package routing
 
 import (
 	"context"
@@ -42,15 +42,6 @@ func DefaultCrawlConfig() *CrawlConfig {
 	}
 }
 
-type crawlJob[K kad.Key[K], N kad.NodeID[K]] struct {
-	node   N
-	target K
-}
-
-func (c *crawlJob[K, N]) mapKey() string {
-	return c.node.String() + key.HexString(c.target)
-}
-
 type Crawl[K kad.Key[K], N kad.NodeID[K], M coordt.Message] struct {
 	self N
 	id   coordt.QueryID
@@ -74,10 +65,6 @@ func NewCrawl[K kad.Key[K], N kad.NodeID[K], M coordt.Message](self N, id coordt
 		return nil, err
 	}
 
-	if len(seed) == 0 {
-		return nil, fmt.Errorf("empty seed")
-	}
-
 	c := &Crawl[K, N, M]{
 		self:    self,
 		id:      id,
@@ -113,6 +100,10 @@ func NewCrawl[K kad.Key[K], N kad.NodeID[K], M coordt.Message](self N, id coordt
 		}
 	}
 
+	if len(seed) == 0 {
+		return nil, fmt.Errorf("empty seed")
+	}
+
 	return c, nil
 }
 
@@ -148,18 +139,18 @@ func (c *Crawl[K, N, M]) Advance(ctx context.Context, ev CrawlEvent) (out CrawlS
 					continue
 				}
 
-				job := crawlJob[K, N]{
+				newJob := crawlJob[K, N]{
 					node:   node,
 					target: target.Key(),
 				}
 
-				mapKey := job.mapKey()
+				newMapKey := newJob.mapKey()
 
-				if _, found := c.cpls[mapKey]; found {
+				if _, found := c.cpls[newMapKey]; found {
 					continue
 				}
 
-				c.cpls[mapKey] = i
+				c.cpls[newMapKey] = i
 				c.todo = append(c.todo, job)
 			}
 		}
@@ -216,6 +207,15 @@ func (c *Crawl[K, N, M]) Advance(ctx context.Context, ev CrawlEvent) (out CrawlS
 	return &StateCrawlFinished{}
 }
 
+type crawlJob[K kad.Key[K], N kad.NodeID[K]] struct {
+	node   N
+	target K
+}
+
+func (c *crawlJob[K, N]) mapKey() string {
+	return c.node.String() + key.HexString(c.target)
+}
+
 type CrawlState interface {
 	crawlState()
 }

diff --git a/internal/coord/query/crawl_test.go → internal/coord/routing/crawl_test.go b/internal/coord/query/crawl_test.go → internal/coord/routing/crawl_test.go
@@ -1,4 +1,4 @@
-package query
+package routing
 
 import (
 	"context"
@@ -14,6 +14,40 @@ import (
 
 var _ coordt.StateMachine[CrawlEvent, CrawlState] = (*Crawl[tiny.Key, tiny.Node, tiny.Message])(nil)
 
+func TestNewCrawl(t *testing.T) {
+	self := tiny.NewNode(0)
+	a := tiny.NewNode(0b10000100)
+	b := tiny.NewNode(0b11000000)
+
+	t.Run("initializes maps", func(t *testing.T) {
+		cfg := DefaultCrawlConfig()
+		cfg.MaxCPL = 4
+		seed := []tiny.Node{a}
+		qry, err := NewCrawl[tiny.Key, tiny.Node, tiny.Message](self, coordt.QueryID("test"), tiny.NodeWithCpl, seed, cfg)
+		require.NoError(t, err)
+		require.NotNil(t, qry)
+		require.Len(t, qry.todo, 4)
+		require.NotNil(t, qry.waiting)
+		require.NotNil(t, qry.success)
+		require.NotNil(t, qry.failed)
+		require.NotNil(t, qry.errors)
+	})
+
+	t.Run("removes self from seed", func(t *testing.T) {
+		cfg := DefaultCrawlConfig()
+		cfg.MaxCPL = 4
+		seed := []tiny.Node{self, a, b}
+		qry, err := NewCrawl[tiny.Key, tiny.Node, tiny.Message](self, coordt.QueryID("test"), tiny.NodeWithCpl, seed, cfg)
+		require.NoError(t, err)
+		require.NotNil(t, qry)
+		require.Len(t, qry.todo, cfg.MaxCPL*2) // self is not included
+		require.NotNil(t, qry.waiting)
+		require.NotNil(t, qry.success)
+		require.NotNil(t, qry.failed)
+		require.NotNil(t, qry.errors)
+	})
+}
+
 func TestCrawl_Advance(t *testing.T) {
 	ctx := context.Background()
 
@@ -87,13 +121,29 @@ func TestCrawl_Advance(t *testing.T) {
 	assert.Len(t, qry.success, 2)
 	assert.Len(t, qry.failed, 0)
 
+	moreReqs := make([]*StateCrawlFindCloser[tiny.Key, tiny.Node], cfg.MaxCPL)
+	moreReqs[0] = tstate
+	for i := 1; i < cfg.MaxCPL; i++ {
+		state = qry.Advance(ctx, &EventCrawlPoll{})
+		tstate, ok = state.(*StateCrawlFindCloser[tiny.Key, tiny.Node])
+		require.True(t, ok, "type is %T", state)
+		moreReqs[i] = tstate
+	}
+
 	for i := 2; i < len(reqs); i++ {
 		state = qry.Advance(ctx, &EventCrawlNodeResponse[tiny.Key, tiny.Node]{
 			NodeID:      reqs[i].NodeID,
 			Target:      reqs[i].Target,
 			CloserNodes: []tiny.Node{},
 		})
 	}
+	for i := 0; i < len(moreReqs); i++ {
+		state = qry.Advance(ctx, &EventCrawlNodeResponse[tiny.Key, tiny.Node]{
+			NodeID:      moreReqs[i].NodeID,
+			Target:      moreReqs[i].Target,
+			CloserNodes: []tiny.Node{},
+		})
+	}
 
 	require.IsType(t, &StateCrawlIdle{}, state)
 }