use a notification type strategy for the query events

core/commands/dht.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	cmds "github.com/jbenet/go-ipfs/commands"
+	notif "github.com/jbenet/go-ipfs/notifications"
 	ipdht "github.com/jbenet/go-ipfs/routing/dht"
 	u "github.com/jbenet/go-ipfs/util"
 )
@@ -46,16 +47,18 @@ var queryDhtCmd = &cmds.Command{
 			return nil, errors.New("Routing service was not a dht")
 		}
 
-		events := make(chan *ipdht.QueryEvent)
+		events := make(chan *notif.QueryEvent)
-		closestPeers, err := dht.GetClosestPeers(req.Context().Context, u.Key(req.Arguments()[0]), events)
+		ctx := notif.RegisterForQueryEvents(req.Context().Context, events)
+
+		closestPeers, err := dht.GetClosestPeers(ctx, u.Key(req.Arguments()[0]))
 
 		go func() {
 			defer close(events)
 			for p := range closestPeers {
-				events <- &ipdht.QueryEvent{
+				notif.PublishQueryEvent(ctx, &notif.QueryEvent{
 					ID:   p,
-					Type: ipdht.FinalPeer,
+					Type: notif.FinalPeer,
-				}
+				})
 			}
 		}()
 
@@ -76,7 +79,7 @@ var queryDhtCmd = &cmds.Command{
 			}
 
 			marshal := func(v interface{}) (io.Reader, error) {
-				obj, ok := v.(*ipdht.QueryEvent)
+				obj, ok := v.(*notif.QueryEvent)
 				if !ok {
 					return nil, u.ErrCast()
 				}
@@ -84,15 +87,15 @@ var queryDhtCmd = &cmds.Command{
 				buf := new(bytes.Buffer)
 				fmt.Fprintf(buf, "%s: ", time.Now().Format("15:04:05.000"))
 				switch obj.Type {
-				case ipdht.FinalPeer:
+				case notif.FinalPeer:
 					fmt.Fprintf(buf, "%s\n", obj.ID)
-				case ipdht.PeerResponse:
+				case notif.PeerResponse:
 					fmt.Fprintf(buf, "* %s says use ", obj.ID)
 					for _, p := range obj.Responses {
 						fmt.Fprintf(buf, "%s ", p.ID)
 					}
 					fmt.Fprintln(buf)
-				case ipdht.SendingQuery:
+				case notif.SendingQuery:
 					fmt.Fprintf(buf, "* querying %s\n", obj.ID)
 				}
 				return buf, nil
@@ -104,5 +107,5 @@ var queryDhtCmd = &cmds.Command{
 			}, nil
 		},
 	},
-	Type: ipdht.QueryEvent{},
+	Type: notif.QueryEvent{},
 }

notifications/query.go

@@ -0,0 +1,73 @@
+package notifications
+
+import (
+	"encoding/json"
+
+	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
+	peer "github.com/jbenet/go-ipfs/p2p/peer"
+)
+
+const RoutingQueryKey = "RoutingQueryEvent"
+
+type QueryEventType int
+
+const (
+	SendingQuery QueryEventType = iota
+	PeerResponse
+	FinalPeer
+)
+
+type QueryEvent struct {
+	ID        peer.ID
+	Type      QueryEventType
+	Responses []*peer.PeerInfo
+}
+
+func RegisterForQueryEvents(ctx context.Context, ch chan<- *QueryEvent) context.Context {
+	return context.WithValue(ctx, RoutingQueryKey, ch)
+}
+
+func PublishQueryEvent(ctx context.Context, ev *QueryEvent) {
+	ich := ctx.Value(RoutingQueryKey)
+	if ich == nil {
+		return
+	}
+
+	ch, ok := ich.(chan<- *QueryEvent)
+	if !ok {
+		return
+	}
+
+	select {
+	case ch <- ev:
+	case <-ctx.Done():
+	}
+}
+
+func (qe *QueryEvent) MarshalJSON() ([]byte, error) {
+	out := make(map[string]interface{})
+	out["ID"] = peer.IDB58Encode(qe.ID)
+	out["Type"] = int(qe.Type)
+	out["Responses"] = qe.Responses
+	return json.Marshal(out)
+}
+
+func (qe *QueryEvent) UnmarshalJSON(b []byte) error {
+	temp := struct {
+		ID        string
+		Type      int
+		Responses []*peer.PeerInfo
+	}{}
+	err := json.Unmarshal(b, &temp)
+	if err != nil {
+		return err
+	}
+	pid, err := peer.IDB58Decode(temp.ID)
+	if err != nil {
+		return err
+	}
+	qe.ID = pid
+	qe.Type = QueryEventType(temp.Type)
+	qe.Responses = temp.Responses
+	return nil
+}

routing/dht/lookup.go

@@ -1,10 +1,9 @@
 package dht
 
 import (
-	"encoding/json"
-
 	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
 
+	notif "github.com/jbenet/go-ipfs/notifications"
 	peer "github.com/jbenet/go-ipfs/p2p/peer"
 	kb "github.com/jbenet/go-ipfs/routing/kbucket"
 	u "github.com/jbenet/go-ipfs/util"
@@ -12,20 +11,7 @@ import (
 	pset "github.com/jbenet/go-ipfs/util/peerset"
 )
 
-type QueryEventType int
+// Required in order for proper JSON marshaling
-
-const (
-	SendingQuery QueryEventType = iota
-	PeerResponse
-	FinalPeer
-)
-
-type QueryEvent struct {
-	ID        peer.ID
-	Type      QueryEventType
-	Responses []*peer.PeerInfo
-}
-
 func pointerizePeerInfos(pis []peer.PeerInfo) []*peer.PeerInfo {
 	out := make([]*peer.PeerInfo, len(pis))
 	for i, p := range pis {
@@ -37,7 +23,7 @@ func pointerizePeerInfos(pis []peer.PeerInfo) []*peer.PeerInfo {
 
 // Kademlia 'node lookup' operation. Returns a channel of the K closest peers
 // to the given key
-func (dht *IpfsDHT) GetClosestPeers(ctx context.Context, key u.Key, events chan<- *QueryEvent) (<-chan peer.ID, error) {
+func (dht *IpfsDHT) GetClosestPeers(ctx context.Context, key u.Key) (<-chan peer.ID, error) {
 	e := log.EventBegin(ctx, "getClosestPeers", &key)
 	tablepeers := dht.routingTable.NearestPeers(kb.ConvertKey(key), AlphaValue)
 	if len(tablepeers) == 0 {
@@ -58,12 +44,10 @@ func (dht *IpfsDHT) GetClosestPeers(ctx context.Context, key u.Key, events chan<
 
 	query := dht.newQuery(key, func(ctx context.Context, p peer.ID) (*dhtQueryResult, error) {
 		// For DHT query command
-		select {
+		notif.PublishQueryEvent(ctx, &notif.QueryEvent{
-		case events <- &QueryEvent{
+			Type: notif.SendingQuery,
-			Type: SendingQuery,
 			ID:   p,
-		}:
+		})
-		}
 
 		closer, err := dht.closerPeersSingle(ctx, key, p)
 		if err != nil {
@@ -86,13 +70,11 @@ func (dht *IpfsDHT) GetClosestPeers(ctx context.Context, key u.Key, events chan<
 		log.Errorf("filtered: %v", filtered)
 
 		// For DHT query command
-		select {
+		notif.PublishQueryEvent(ctx, &notif.QueryEvent{
-		case events <- &QueryEvent{
+			Type:      notif.PeerResponse,
-			Type:      PeerResponse,
 			ID:        p,
 			Responses: pointerizePeerInfos(filtered),
-		}:
+		})
-		}
 
 		return &dhtQueryResult{closerPeers: filtered}, nil
 	})
@@ -126,31 +108,3 @@ func (dht *IpfsDHT) closerPeersSingle(ctx context.Context, key u.Key, p peer.ID)
 	}
 	return out, nil
 }
-
-func (qe *QueryEvent) MarshalJSON() ([]byte, error) {
-	out := make(map[string]interface{})
-	out["ID"] = peer.IDB58Encode(qe.ID)
-	out["Type"] = int(qe.Type)
-	out["Responses"] = qe.Responses
-	return json.Marshal(out)
-}
-
-func (qe *QueryEvent) UnmarshalJSON(b []byte) error {
-	temp := struct {
-		ID        string
-		Type      int
-		Responses []*peer.PeerInfo
-	}{}
-	err := json.Unmarshal(b, &temp)
-	if err != nil {
-		return err
-	}
-	pid, err := peer.IDB58Decode(temp.ID)
-	if err != nil {
-		return err
-	}
-	qe.ID = pid
-	qe.Type = QueryEventType(temp.Type)
-	qe.Responses = temp.Responses
-	return nil
-}

routing/dht/routing.go

@@ -48,7 +48,7 @@ func (dht *IpfsDHT) PutValue(ctx context.Context, key u.Key, value []byte) error
 		return err
 	}
 
-	pchan, err := dht.GetClosestPeers(ctx, key, nil)
+	pchan, err := dht.GetClosestPeers(ctx, key)
 	if err != nil {
 		return err
 	}
@@ -134,7 +134,7 @@ func (dht *IpfsDHT) Provide(ctx context.Context, key u.Key) error {
 	// add self locally
 	dht.providers.AddProvider(key, dht.self)
 
-	peers, err := dht.GetClosestPeers(ctx, key, nil)
+	peers, err := dht.GetClosestPeers(ctx, key)
 	if err != nil {
 		return err
 	}