feat_: hash based query for outgoing messages.

Author: kaichaosun

Diff for: protocol/messenger_peers.go

@@ -13,7 +13,7 @@ func (m *Messenger) AddStorePeer(address string) (peer.ID, error) {
 }
 
 func (m *Messenger) AddRelayPeer(address string) (peer.ID, error) {
-	return m.transport.AddStorePeer(address)
+	return m.transport.AddRelayPeer(address)
 }
 
 func (m *Messenger) DialPeer(address string) error {

Diff for: wakuv2/common/message.go

@@ -22,6 +22,7 @@ type MessageType int
 const (
 	RelayedMessageType MessageType = iota
 	StoreMessageType
+	SendMessageType
 )
 
 // MessageParams specifies the exact way a message should be wrapped
@@ -46,7 +47,7 @@ type ReceivedMessage struct {
 	Padding   []byte
 	Signature []byte
 
-	Sent uint32           // Time when the message was posted into the network
+	Sent uint32           // Time when the message was posted into the network in seconds
 	Src  *ecdsa.PublicKey // Message recipient (identity used to decode the message)
 	Dst  *ecdsa.PublicKey // Message recipient (identity used to decode the message)
 

Diff for: wakuv2/waku.go

@@ -19,6 +19,7 @@
 package wakuv2
 
 import (
+	"bytes"
 	"context"
 	"crypto/ecdsa"
 	"crypto/sha256"
@@ -65,6 +66,7 @@ import (
 	"github.com/waku-org/go-waku/waku/v2/protocol/lightpush"
 	"github.com/waku-org/go-waku/waku/v2/protocol/peer_exchange"
 	"github.com/waku-org/go-waku/waku/v2/protocol/relay"
+	"github.com/waku-org/go-waku/waku/v2/protocol/store"
 
 	"github.com/status-im/status-go/connection"
 	"github.com/status-im/status-go/eth-node/types"
@@ -126,6 +128,9 @@ type Waku struct {
 	storeMsgIDs   map[gethcommon.Hash]bool // Map of the currently processing ids
 	storeMsgIDsMu sync.RWMutex
 
+	sendMsgIDs   map[string]map[gethcommon.Hash]uint32
+	sendMsgIDsMu sync.RWMutex
+
 	topicHealthStatusChan   chan peermanager.TopicHealthStatus
 	connStatusSubscriptions map[string]*types.ConnStatusSubscription
 	connStatusMu            sync.Mutex
@@ -206,6 +211,8 @@ func New(nodeKey *ecdsa.PrivateKey, fleet string, cfg *Config, logger *zap.Logge
 		storeMsgIDs:                     make(map[gethcommon.Hash]bool),
 		timesource:                      ts,
 		storeMsgIDsMu:                   sync.RWMutex{},
+		sendMsgIDs:                      make(map[string]map[gethcommon.Hash]uint32),
+		sendMsgIDsMu:                    sync.RWMutex{},
 		logger:                          logger,
 		discV5BootstrapNodes:            cfg.DiscV5BootstrapNodes,
 		onHistoricMessagesRequestFailed: onHistoricMessagesRequestFailed,
@@ -970,6 +977,38 @@ func (w *Waku) broadcast() {
 	}
 }
 
+func (w *Waku) checkIfMessagesStored() {
+	ticker := time.NewTicker(time.Second * 5)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-w.ctx.Done():
+			w.logger.Debug("Stop the look for message stored check")
+			return
+		case <-ticker.C:
+			w.logger.Debug("Running loop for messages stored check")
+			w.logger.Debug("Send Message IDs", zap.Any("sendMsgIDs", w.sendMsgIDs))
+			w.sendMsgIDsMu.Lock()
+			for pubsubTopic, subMsgs := range w.sendMsgIDs {
+				var queryMsgIds []gethcommon.Hash
+				for msgID, sendTime := range subMsgs {
+					// message is sent 5 seconds ago, check if it's stored
+					if uint32(w.timesource.Now().Unix()) > sendTime+5 {
+						queryMsgIds = append(queryMsgIds, msgID)
+					}
+				}
+				w.logger.Debug("Store query for message hashes", zap.Any("queryMsgIds", queryMsgIds), zap.String("pubsubTopic", pubsubTopic))
+				if len(queryMsgIds) > 0 {
+					w.messageHashBasedQuery(w.ctx, queryMsgIds, pubsubTopic)
+				}
+			}
+
+			w.sendMsgIDsMu.Unlock()
+		}
+	}
+}
+
 type publishFn = func(envelope *protocol.Envelope, logger *zap.Logger) error
 
 func (w *Waku) publishEnvelope(envelope *protocol.Envelope, publishFn publishFn, logger *zap.Logger) {
@@ -979,14 +1018,11 @@ func (w *Waku) publishEnvelope(envelope *protocol.Envelope, publishFn publishFn,
 	if err := publishFn(envelope, logger); err != nil {
 		logger.Error("could not send message", zap.Error(err))
 		event = common.EventEnvelopeExpired
-	} else {
-		event = common.EventEnvelopeSent
+		w.SendEnvelopeEvent(common.EnvelopeEvent{
+			Hash:  gethcommon.BytesToHash(envelope.Hash().Bytes()),
+			Event: event,
+		})
 	}
-
-	w.SendEnvelopeEvent(common.EnvelopeEvent{
-		Hash:  gethcommon.BytesToHash(envelope.Hash().Bytes()),
-		Event: event,
-	})
 }
 
 // Send injects a message into the waku send queue, to be distributed in the
@@ -1015,14 +1051,88 @@ func (w *Waku) Send(pubsubTopic string, msg *pb.WakuMessage) ([]byte, error) {
 	alreadyCached := w.envelopeCache.Has(gethcommon.BytesToHash(envelope.Hash().Bytes()))
 	w.poolMu.Unlock()
 	if !alreadyCached {
-		recvMessage := common.NewReceivedMessage(envelope, common.RelayedMessageType)
+		recvMessage := common.NewReceivedMessage(envelope, common.SendMessageType)
 		w.postEvent(recvMessage) // notify the local node about the new message
 		w.addEnvelope(recvMessage)
 	}
 
 	return envelope.Hash().Bytes(), nil
 }
 
+// ctx, peer, r.PubsubTopic, contentTopics, uint64(r.From), uint64(r.To), options, processEnvelopes
+func (w *Waku) messageHashBasedQuery(ctx context.Context, hashes []gethcommon.Hash, pubsubTopic string) {
+	selectedPeers, err := w.node.PeerManager().SelectPeers(
+		peermanager.PeerSelectionCriteria{
+			SelectionType: peermanager.Automatic,
+			Proto:         store.StoreQueryID_v300,
+			PubsubTopics:  []string{pubsubTopic},
+			Ctx:           ctx,
+		},
+	)
+	if err != nil {
+		w.logger.Warn("could not select peers", zap.Error(err))
+		return
+	}
+
+	var opts []store.RequestOption
+	requestID := protocol.GenerateRequestID()
+	opts = append(opts, store.WithRequestID(requestID))
+	opts = append(opts, store.WithPeer(selectedPeers[0]))
+
+	messageHashes := make([]pb.MessageHash, len(hashes))
+	for i, hash := range hashes {
+		messageHashes[i] = pb.ToMessageHash(hash.Bytes())
+	}
+
+	result, err := w.node.Store().QueryByHash(ctx, messageHashes, opts...)
+	if err != nil {
+		w.logger.Warn("store.queryByHash failed", zap.String("requestID", hexutil.Encode(requestID)), zap.Error(err))
+		return
+	}
+
+	var ackHashes []gethcommon.Hash
+	var missedHashes []gethcommon.Hash
+	for _, hash := range hashes {
+		found := false
+		for _, msg := range result.Messages() {
+			if bytes.Equal(msg.GetMessageHash(), hash.Bytes()) {
+				found = true
+				break
+			}
+		}
+		if found {
+			ackHashes = append(ackHashes, hash)
+		} else {
+			missedHashes = append(missedHashes, hash)
+		}
+
+		subMsgs := w.sendMsgIDs[pubsubTopic]
+		delete(subMsgs, hash)
+		if len(subMsgs) == 0 {
+			delete(w.sendMsgIDs, pubsubTopic)
+		} else {
+			w.sendMsgIDs[pubsubTopic] = subMsgs
+		}
+	}
+
+	w.logger.Debug("Ack message hashes", zap.Any("ackHashes", ackHashes))
+	w.logger.Debug("Missed message hashes", zap.Any("missedHashes", missedHashes))
+
+	for _, hash := range ackHashes {
+		w.SendEnvelopeEvent(common.EnvelopeEvent{
+			Hash:  hash,
+			Event: common.EventEnvelopeSent,
+		})
+	}
+
+	for _, hash := range missedHashes {
+		w.SendEnvelopeEvent(common.EnvelopeEvent{
+			Hash:  hash,
+			Event: common.EventEnvelopeExpired,
+		})
+	}
+}
+
 func (w *Waku) query(ctx context.Context, peerID peer.ID, pubsubTopic string, topics []common.TopicType, from uint64, to uint64, requestID []byte, opts []legacy_store.HistoryRequestOption) (*legacy_store.Result, error) {
 
 	if len(requestID) != 0 {
@@ -1203,6 +1313,8 @@ func (w *Waku) Start() error {
 
 	go w.broadcast()
 
+	go w.checkIfMessagesStored()
+
 	// we should wait `seedBootnodesForDiscV5` shutdown smoothly before set w.ctx to nil within `w.Stop()`
 	w.wg.Add(1)
 	go w.seedBootnodesForDiscV5()
@@ -1382,6 +1494,17 @@ func (w *Waku) processReceivedMessage(e *common.ReceivedMessage) {
 		w.storeMsgIDsMu.Unlock()
 	}
 
+	if e.MsgType == common.SendMessageType && !(*e.Envelope.Message().Ephemeral) {
+		w.sendMsgIDsMu.Lock()
+		subMsgs, ok := w.sendMsgIDs[e.PubsubTopic]
+		if !ok {
+			subMsgs = make(map[gethcommon.Hash]uint32)
+		}
+		subMsgs[e.Hash()] = e.Sent
+		w.sendMsgIDs[e.PubsubTopic] = subMsgs
+		w.sendMsgIDsMu.Unlock()
+	}
+
 	matched := w.filters.NotifyWatchers(e)
 
 	// If not matched we remove it