Improve query capabilities

v2/coord/coordinator.go

@@ -9,6 +9,7 @@
 	"time"
 
 	"github.com/benbjohnson/clock"
+	uuid "github.com/google/uuid"
 	logging "github.com/ipfs/go-log/v2"
 	"github.com/plprobelab/go-kademlia/kad"
 	"github.com/plprobelab/go-kademlia/kaderr"
@@ -167,7 +168,7 @@
 
 func NewCoordinator(self kadt.PeerID, rtr Router[kadt.Key, kadt.PeerID, *pb.Message], rt routing.RoutingTableCpl[kadt.Key, kadt.PeerID], cfg *CoordinatorConfig) (*Coordinator, error) {
 	if cfg == nil {
 		cfg = DefaultCoordinatorConfig()
 	} else if err := cfg.Validate(); err != nil {
 		return nil, err
 	}
@@ -175,8 +176,8 @@
 	// initialize a new telemetry struct
 	tele, err := NewTelemetry(cfg.MeterProvider, cfg.TracerProvider)
 	if err != nil {
 		return nil, fmt.Errorf("init telemetry: %w", err)
 	}
 
 	qpCfg := query.DefaultPoolConfig()
 	qpCfg.Clock = cfg.Clock
@@ -185,7 +186,7 @@
 	qpCfg.QueryConcurrency = cfg.RequestConcurrency
 	qpCfg.RequestTimeout = cfg.RequestTimeout
 
-	qp, err := query.NewPool[kadt.Key](self, qpCfg)
+	qp, err := query.NewPool[kadt.Key, kadt.PeerID, *pb.Message](self, qpCfg)
 	if err != nil {
 		return nil, fmt.Errorf("query pool: %w", err)
 	}
@@ -197,7 +198,7 @@
 	bootstrapCfg.RequestConcurrency = cfg.RequestConcurrency
 	bootstrapCfg.RequestTimeout = cfg.RequestTimeout
 
-	bootstrap, err := routing.NewBootstrap[kadt.Key](kadt.PeerID(self), bootstrapCfg)
+	bootstrap, err := routing.NewBootstrap(kadt.PeerID(self), bootstrapCfg)
 	if err != nil {
 		return nil, fmt.Errorf("bootstrap: %w", err)
 	}
@@ -259,7 +260,7 @@
 	return nil
 }
 
 func (c *Coordinator) ID() kadt.PeerID {
 	return c.self
 }
 
@@ -341,7 +342,7 @@
 
 // GetValue requests that the node return any value associated with the supplied key.
 // If the node does not have a value for the key it returns ErrValueNotFound.
 func (c *Coordinator) GetValue(ctx context.Context, k kadt.Key) (Value, error) {
 	panic("not implemented")
 }
 
@@ -351,8 +352,18 @@
 	panic("not implemented")
 }
 
-// Query traverses the DHT calling fn for each node visited.
-func (c *Coordinator) Query(ctx context.Context, target kadt.Key, fn QueryFunc) (QueryStats, error) {
+// QueryClosest starts a query that attempts to find the closest nodes to the target key.
+// It returns the closest nodes found to the target key and statistics on the actions of the query.
+//
+// The supplied [QueryFunc] is called after each successful request to a node with the ID of the node,
+// the response received from the find nodes request made to the node and the current query stats. The query
+// terminates when [QueryFunc] returns an error or when the query has visited the configured minimum number
+// of closest nodes (default 20)
+//
+// numResults specifies the minimum number of nodes to successfully contact before considering iteration complete.
+// The query is considered to be exhausted when it has received responses from at least this number of nodes
+// and there are no closer nodes remaining to be contacted. A default of 20 is used if this value is less than 1.
+func (c *Coordinator) QueryClosest(ctx context.Context, target kadt.Key, fn QueryFunc, numResults int) ([]kadt.PeerID, QueryStats, error) {
 	ctx, span := c.tele.Tracer.Start(ctx, "Coordinator.Query")
 	defer span.End()
 
@@ -361,7 +372,7 @@
 
 	seeds, err := c.GetClosestNodes(ctx, target, 20)
 	if err != nil {
-		return QueryStats{}, err
+		return nil, QueryStats{}, err
 	}
 
 	seedIDs := make([]kadt.PeerID, 0, len(seeds))
@@ -370,23 +381,75 @@
 	}
 
 	waiter := NewWaiter[BehaviourEvent]()
-	queryID := query.QueryID("foo") // TODO: choose query ID
+	queryID := query.QueryID(uuid.New().String())
 
-	cmd := &EventStartQuery{
+	cmd := &EventStartFindCloserQuery{
 		QueryID:           queryID,
 		Target:            target,
 		KnownClosestNodes: seedIDs,
 		Notify:            waiter,
+		NumResults:        numResults,
 	}
 
 	// queue the start of the query
 	c.queryBehaviour.Notify(ctx, cmd)
 
+	return c.waitForQuery(ctx, queryID, waiter, fn)
+}
+
+// QueryMessage starts a query that iterates over the closest nodes to the target key in the supplied message.
+// The message is sent to each node that is visited.
+//
+// The supplied [QueryFunc] is called after each successful request to a node with the ID of the node,
+// the response received from the find nodes request made to the node and the current query stats. The query
+// terminates when [QueryFunc] returns an error or when the query has visited the configured minimum number
+// of closest nodes (default 20)
+//
+// numResults specifies the minimum number of nodes to successfully contact before considering iteration complete.
+// The query is considered to be exhausted when it has received responses from at least this number of nodes
+// and there are no closer nodes remaining to be contacted. A default of 20 is used if this value is less than 1.
+func (c *Coordinator) QueryMessage(ctx context.Context, msg *pb.Message, fn QueryFunc, numResults int) (QueryStats, error) {
+	ctx, span := c.tele.Tracer.Start(ctx, "Coordinator.QueryMessage")
+	defer span.End()
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	seeds, err := c.GetClosestNodes(ctx, msg.Target(), 20)
+	if err != nil {
+		return QueryStats{}, err
+	}
+
+	seedIDs := make([]kadt.PeerID, 0, len(seeds))
+	for _, s := range seeds {
+		seedIDs = append(seedIDs, kadt.PeerID(s.ID()))
+	}
+
+	waiter := NewWaiter[BehaviourEvent]()
+	queryID := query.QueryID(uuid.New().String())
+
+	cmd := &EventStartMessageQuery{
+		QueryID:           queryID,
+		Target:            msg.Target(),
+		Message:           msg,
+		KnownClosestNodes: seedIDs,
+		Notify:            waiter,
+		NumResults:        numResults,
+	}
+
+	// queue the start of the query
+	c.queryBehaviour.Notify(ctx, cmd)
+
+	_, stats, err := c.waitForQuery(ctx, queryID, waiter, fn)
+	return stats, err
+}
+
+func (c *Coordinator) waitForQuery(ctx context.Context, queryID query.QueryID, waiter *Waiter[BehaviourEvent], fn QueryFunc) ([]kadt.PeerID, QueryStats, error) {
 	var lastStats QueryStats
 	for {
 		select {
 		case <-ctx.Done():
-			return lastStats, ctx.Err()
+			return nil, lastStats, ctx.Err()
 		case wev := <-waiter.Chan():
 			ctx, ev := wev.Ctx, wev.Event
 			switch ev := ev.(type) {
@@ -403,26 +466,22 @@
 					break
 				}
 
-				err = fn(ctx, nh, lastStats)
+				err = fn(ctx, nh.ID(), ev.Response, lastStats)
 				if errors.Is(err, ErrSkipRemaining) {
 					// done
 					c.queryBehaviour.Notify(ctx, &EventStopQuery{QueryID: queryID})
-					return lastStats, nil
-				}
-				if errors.Is(err, ErrSkipNode) {
-					// TODO: don't add closer nodes from this node
-					break
+					return nil, lastStats, nil
 				}
 				if err != nil {
 					// user defined error that terminates the query
 					c.queryBehaviour.Notify(ctx, &EventStopQuery{QueryID: queryID})
-					return lastStats, err
+					return nil, lastStats, err
 				}
 
 			case *EventQueryFinished:
 				// query is done
 				lastStats.Exhausted = true
-				return lastStats, nil
+				return ev.ClosestNodes, lastStats, nil
 
 			default:
 				panic(fmt.Sprintf("unexpected event: %T", ev))
@@ -479,15 +538,15 @@
 
 // NotifyNonConnectivity notifies the coordinator that a peer has failed a connectivity check
 // which means it is not connected and/or it doesn't support finding closer nodes
 func (c *Coordinator) NotifyNonConnectivity(ctx context.Context, id kadt.PeerID) error {
 	ctx, span := c.tele.Tracer.Start(ctx, "Coordinator.NotifyNonConnectivity")
 	defer span.End()
 
 	c.routingBehaviour.Notify(ctx, &EventNotifyNonConnectivity{
 		NodeID: id,
 	})
 
 	return nil
 }
 
 // A BufferedRoutingNotifier is a [RoutingNotifier] that buffers [RoutingNotification] events and provides methods
@@ -530,8 +589,8 @@
 
 		// wait to be signaled that there is a new event
 		select {
 		case <-ctx.Done():
 			return nil, fmt.Errorf("test deadline exceeded while waiting for event %T", expected)
 		case <-w.signal:
 		}
 	}
@@ -556,8 +615,8 @@
 
 		// wait to be signaled that there is a new event
 		select {
 		case <-ctx.Done():
 			return nil, fmt.Errorf("test deadline exceeded while waiting for routing updated event")
 		case <-w.signal:
 		}
 	}
@@ -582,8 +641,8 @@
 
 		// wait to be signaled that there is a new event
 		select {
 		case <-ctx.Done():
 			return nil, fmt.Errorf("test deadline exceeded while waiting for routing removed event")
 		case <-w.signal:
 		}
 	}

v2/coord/coordinator_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/libp2p/go-libp2p-kad-dht/v2/coord/internal/nettest"
 	"github.com/libp2p/go-libp2p-kad-dht/v2/internal/kadtest"
 	"github.com/libp2p/go-libp2p-kad-dht/v2/kadt"
+	"github.com/libp2p/go-libp2p-kad-dht/v2/pb"
 )
 
 const peerstoreTTL = 10 * time.Minute
@@ -115,13 +116,13 @@ func TestExhaustiveQuery(t *testing.T) {
 	visited := make(map[string]int)
 
 	// Record the nodes as they are visited
-	qfn := func(ctx context.Context, node Node, stats QueryStats) error {
-		visited[node.ID().String()]++
+	qfn := func(ctx context.Context, id kadt.PeerID, msg *pb.Message, stats QueryStats) error {
+		visited[id.String()]++
 		return nil
 	}
 
 	// Run a query to find the value
-	_, err = c.Query(ctx, target, qfn)
+	_, _, err = c.QueryClosest(ctx, target, qfn, 20)
 	require.NoError(t, err)
 
 	require.Equal(t, 3, len(visited))
@@ -154,13 +155,13 @@ func TestRoutingUpdatedEventEmittedForCloserNodes(t *testing.T) {
 		log.Fatalf("unexpected error creating coordinator: %v", err)
 	}
 
-	qfn := func(ctx context.Context, node Node, stats QueryStats) error {
+	qfn := func(ctx context.Context, id kadt.PeerID, msg *pb.Message, stats QueryStats) error {
 		return nil
 	}
 
 	// Run a query to find the value
 	target := nodes[3].NodeID.Key()
-	_, err = c.Query(ctx, target, qfn)
+	_, _, err = c.QueryClosest(ctx, target, qfn, 20)
 	require.NoError(t, err)
 
 	// the query run by the dht should have received a response from nodes[1] with closer nodes

v2/coord/coretypes.go

@@ -8,6 +8,7 @@ import (
 	"github.com/plprobelab/go-kademlia/kad"
 
 	"github.com/libp2p/go-libp2p-kad-dht/v2/kadt"
+	"github.com/libp2p/go-libp2p-kad-dht/v2/pb"
 )
 
 // Value is a value that may be stored in the DHT.
@@ -49,7 +50,7 @@ var (
 // Query stops entirely and returns that error.
 //
 // The stats argument contains statistics on the progress of the query so far.
-type QueryFunc func(ctx context.Context, node Node, stats QueryStats) error
+type QueryFunc func(ctx context.Context, id kadt.PeerID, resp *pb.Message, stats QueryStats) error
 
 type QueryStats struct {
 	Start     time.Time // Start is the time the query began executing.

v2/coord/event.go

@@ -3,6 +3,7 @@
 import (
 	"github.com/libp2p/go-libp2p-kad-dht/v2/coord/query"
 	"github.com/libp2p/go-libp2p-kad-dht/v2/kadt"
+	"github.com/libp2p/go-libp2p-kad-dht/v2/pb"
 )
 
 type BehaviourEvent interface {
@@ -60,15 +61,39 @@
 func (*EventOutboundGetCloserNodes) nodeHandlerRequest() {}
 func (*EventOutboundGetCloserNodes) networkCommand()     {}
 
-type EventStartQuery struct {
+type EventOutboundSendMessage struct {
+	QueryID query.QueryID
+	To      kadt.PeerID
+	Message *pb.Message
+	Notify  Notify[BehaviourEvent]
+}
+
+func (*EventOutboundSendMessage) behaviourEvent()     {}
+func (*EventOutboundSendMessage) nodeHandlerRequest() {}
+func (*EventOutboundSendMessage) networkCommand()     {}
+
+type EventStartMessageQuery struct {
+	QueryID           query.QueryID
+	Target            kadt.Key
+	Message           *pb.Message
+	KnownClosestNodes []kadt.PeerID
+	Notify            NotifyCloser[BehaviourEvent]
+	NumResults        int // the minimum number of nodes to successfully contact before considering iteration complete
+}
+
+func (*EventStartMessageQuery) behaviourEvent() {}
+func (*EventStartMessageQuery) queryCommand()   {}
+
+type EventStartFindCloserQuery struct {
 	QueryID           query.QueryID
 	Target            kadt.Key
 	KnownClosestNodes []kadt.PeerID
 	Notify            NotifyCloser[BehaviourEvent]
+	NumResults        int // the minimum number of nodes to successfully contact before considering iteration complete
 }
 
-func (*EventStartQuery) behaviourEvent() {}
-func (*EventStartQuery) queryCommand()   {}
+func (*EventStartFindCloserQuery) behaviourEvent() {}
+func (*EventStartFindCloserQuery) queryCommand()   {}
 
 type EventStopQuery struct {
 	QueryID query.QueryID
@@ -82,8 +107,8 @@
 	NodeID kadt.PeerID
 }
 
 func (*EventAddNode) behaviourEvent() {}
 func (*EventAddNode) routingCommand() {}
 
 // EventGetCloserNodesSuccess notifies a behaviour that a GetCloserNodes request, initiated by an
 // [EventOutboundGetCloserNodes] event has produced a successful response.
@@ -109,12 +134,36 @@
 func (*EventGetCloserNodesFailure) behaviourEvent()      {}
 func (*EventGetCloserNodesFailure) nodeHandlerResponse() {}
 
+// EventSendMessageSuccess notifies a behaviour that a SendMessage request, initiated by an
+// [EventOutboundSendMessage] event has produced a successful response.
+type EventSendMessageSuccess struct {
+	QueryID     query.QueryID
+	To          kadt.PeerID // To is the peer that the SendMessage request was sent to.
+	Response    *pb.Message
+	CloserNodes []kadt.PeerID
+}
+
+func (*EventSendMessageSuccess) behaviourEvent()      {}
+func (*EventSendMessageSuccess) nodeHandlerResponse() {}
+
+// EventSendMessageFailure notifies a behaviour that a SendMessage request, initiated by an
+// [EventOutboundSendMessage] event has failed to produce a valid response.
+type EventSendMessageFailure struct {
+	QueryID query.QueryID
+	To      kadt.PeerID // To is the peer that the SendMessage request was sent to.
+	Target  kadt.Key
+	Err     error
+}
+
+func (*EventSendMessageFailure) behaviourEvent()      {}
+func (*EventSendMessageFailure) nodeHandlerResponse() {}
+
 // EventQueryProgressed is emitted by the coordinator when a query has received a
 // response from a node.
 type EventQueryProgressed struct {
 	QueryID  query.QueryID
 	NodeID   kadt.PeerID
-	Response Message
+	Response *pb.Message
 	Stats    query.QueryStats
 }
 
@@ -123,8 +172,9 @@
 // EventQueryFinished is emitted by the coordinator when a query has finished, either through
 // running to completion or by being canceled.
 type EventQueryFinished struct {
-	QueryID query.QueryID
-	Stats   query.QueryStats
+	QueryID      query.QueryID
+	Stats        query.QueryStats
+	ClosestNodes []kadt.PeerID
 }
 
 func (*EventQueryFinished) behaviourEvent() {}
@@ -142,8 +192,8 @@
 	NodeID kadt.PeerID
 }
 
 func (*EventRoutingRemoved) behaviourEvent()      {}
 func (*EventRoutingRemoved) routingNotification() {}
 
 // EventBootstrapFinished is emitted by the coordinator when a bootstrap has finished, either through
 // running to completion or by being canceled.
@@ -162,8 +212,8 @@
 	NodeID kadt.PeerID
 }
 
 func (*EventNotifyConnectivity) behaviourEvent()      {}
 func (*EventNotifyConnectivity) routingNotification() {}
 
 // EventNotifyNonConnectivity notifies a behaviour that a peer does not have connectivity and/or does not support
 // finding closer nodes is known.
@@ -171,5 +221,5 @@
 	NodeID kadt.PeerID
 }
 
 func (*EventNotifyNonConnectivity) behaviourEvent() {}
 func (*EventNotifyNonConnectivity) routingCommand() {}

v2/coord/event_test.go

@@ -8,7 +8,8 @@ var (
 )
 
 var (
-	_ QueryCommand = (*EventStartQuery)(nil)
+	_ QueryCommand = (*EventStartMessageQuery)(nil)
+	_ QueryCommand = (*EventStartFindCloserQuery)(nil)
 	_ QueryCommand = (*EventStopQuery)(nil)
 )
 

v2/coord/internal/tiny/node.go

@@ -12,6 +12,10 @@
 	key Key
 }
 
+type Message struct {
+	Content string
+}
+
 var _ kad.NodeID[Key] = Node{}
 
 func NewNode(k Key) Node {
@@ -22,8 +26,8 @@
 	return n.key
 }
 
 func (n Node) Equal(other Node) bool {
 	return n.key.Compare(other.key) == 0
 }
 
 func (n Node) String() string {

v2/coord/network.go

@@ -59,6 +59,16 @@
 		}
 		b.nodeHandlersMu.Unlock()
 		nh.Notify(ctx, ev)
+	case *EventOutboundSendMessage:
+		b.nodeHandlersMu.Lock()
+		p := kadt.PeerID(ev.To)
+		nh, ok := b.nodeHandlers[p]
+		if !ok {
+			nh = NewNodeHandler(p, b.rtr, b.logger, b.tracer)
+			b.nodeHandlers[p] = nh
+		}
+		b.nodeHandlersMu.Unlock()
+		nh.Notify(ctx, ev)
 	default:
 		panic(fmt.Sprintf("unexpected dht event: %T", ev))
 	}
@@ -160,6 +170,26 @@
 			Target:      cmd.Target,
 			CloserNodes: nodes,
 		})
+	case *EventOutboundSendMessage:
+		if cmd.Notify == nil {
+			break
+		}
+		resp, err := h.rtr.SendMessage(ctx, h.self, cmd.Message)
+		if err != nil {
+			cmd.Notify.Notify(ctx, &EventSendMessageFailure{
+				QueryID: cmd.QueryID,
+				To:      h.self,
+				Err:     fmt.Errorf("NodeHandler: %w", err),
+			})
+			return false
+		}
+
+		cmd.Notify.Notify(ctx, &EventSendMessageSuccess{
+			QueryID:     cmd.QueryID,
+			To:          h.self,
+			Response:    resp,
+			CloserNodes: resp.CloserNodes(),
+		})
 	default:
 		panic(fmt.Sprintf("unexpected command type: %T", cmd))
 	}
@@ -215,7 +245,7 @@
 
 // GetValue requests that the node return any value associated with the supplied key.
 // If the node does not have a value for the key it returns ErrValueNotFound.
 func (h *NodeHandler) GetValue(ctx context.Context, key kadt.Key) (Value, error) {
 	panic("not implemented")
 }