Proto tx messages deduplicator

pkg/node/helpers.go

@@ -1,6 +1,11 @@
 package node
 
 import (
+	"context"
+	"hash/maphash"
+	"sync"
+
+	"github.com/wavesplatform/gowaves/pkg/p2p/peer"
 	"github.com/wavesplatform/gowaves/pkg/proto"
 	"github.com/wavesplatform/gowaves/pkg/state"
 	"github.com/wavesplatform/gowaves/pkg/types"
@@ -33,3 +38,176 @@ func maybeEnableExtendedApi(state startProvidingExtendedApi, lastBlock *proto.Bl
 	}
 	return nil
 }
+
+type safeMap[K comparable, V any] struct {
+	mu sync.Mutex
+	m  map[K]V
+}
+
+func newSafeMap[K comparable, V any]() *safeMap[K, V] {
+	return &safeMap[K, V]{
+		m: make(map[K]V),
+	}
+}
+
+func (s *safeMap[K, V]) SetIfNew(key K, value V) bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if _, exists := s.m[key]; exists {
+		return false
+	}
+	s.m[key] = value
+	return true
+}
+
+func (s *safeMap[K, V]) Delete(key K) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	delete(s.m, key)
+}
+
+type protoMessageWrapper struct {
+	protoMsg  peer.ProtoMessage
+	messageID uint64
+}
+
+// messageIDFilterFunc is a function type that takes a proto.Message and returns a uint64
+// representing the deterministic message ID. This is used to filter messages based on
+// their message ID: it can be calculated from the message payload, or from the whole message itself.
+//
+// If the returned message ID is 0, the message bypasses filtering.
+type messageIDFilterFunc func(message proto.Message) uint64
+
+type txMessagePayloadIDFilter struct {
+	seed maphash.Seed
+}
+
+func (f txMessagePayloadIDFilter) Filter(message proto.Message) uint64 {
+	switch msg := message.(type) {
+	case *proto.TransactionMessage:
+		return maphash.Bytes(f.seed, msg.Transaction)
+	case *proto.PBTransactionMessage:
+		return maphash.Bytes(f.seed, msg.Transaction)
+	default:
+		return 0 // non-transaction messages bypass filtering
+	}
+}
+
+func messagesSplitter(
+	ctx context.Context,
+	origMessageCh <-chan peer.ProtoMessage,
+	noFilteredChan, filteredChan chan<- protoMessageWrapper,
+	sm *safeMap[uint64, struct{}],
+	messageIDFilter messageIDFilterFunc,
+) {
+	for {
+		var (
+			msg       peer.ProtoMessage
+			messageID uint64
+		)
+		select {
+		case <-ctx.Done():
+			return
+		case msg = <-origMessageCh:
+			messageID = messageIDFilter(msg.Message)
+		}
+		bypassesFiltering := messageID == 0 // if messageID is 0, it means the message no need to be filtered.
+		needToWrite := bypassesFiltering || sm.SetIfNew(messageID, struct{}{})
+		if !needToWrite {
+			continue // skip a message if message ID already exists in the map
+		}
+		outCh := noFilteredChan
+		if !bypassesFiltering {
+			outCh = filteredChan
+		}
+		select {
+		case <-ctx.Done():
+			return
+		case outCh <- protoMessageWrapper{msg, messageID}:
+		}
+	}
+}
+
+func runMessagesMerger(
+	ctx context.Context, wg *sync.WaitGroup,
+	noFilteredChan, filteredChan <-chan protoMessageWrapper,
+	sm *safeMap[uint64, struct{}],
+) <-chan peer.ProtoMessage {
+	wg.Add(1)
+	outMessageCh := make(chan peer.ProtoMessage) // intentionally unbuffered channel because of map usage
+	go func() {
+		defer wg.Done()
+		for {
+			var wMsg protoMessageWrapper
+			select {
+			case <-ctx.Done():
+				return
+			case wMsg = <-noFilteredChan:
+			case wMsg = <-filteredChan:
+			}
+			select {
+			case <-ctx.Done():
+				return
+			case outMessageCh <- wMsg.protoMsg:
+				if wMsg.messageID != 0 {
+					sm.Delete(wMsg.messageID) // remove message ID from the map after sending the message
+				}
+			}
+		}
+	}()
+	return outMessageCh
+}
+
+type chanLenProvider[T any] <-chan T
+
+func (c chanLenProvider[T]) Len() int { return len(c) }
+
+type aggregatedLenProvider []lenProvider
+
+func (a aggregatedLenProvider) Len() int {
+	l := 0
+	for _, p := range a {
+		l += p.Len()
+	}
+	return l
+}
+
+type waiter interface {
+	Wait()
+}
+
+func deduplicateProtoMessages(
+	ctx context.Context, origMessageCh <-chan peer.ProtoMessage,
+	messageIDFilter messageIDFilterFunc,
+) (<-chan peer.ProtoMessage, lenProvider, waiter) {
+	const noFilteredChanSize = 1000
+
+	wg := &sync.WaitGroup{}
+	sm := newSafeMap[uint64, struct{}]()
+
+	noFilteredChan := make(chan protoMessageWrapper, noFilteredChanSize)
+	filteredChan := make(chan protoMessageWrapper, max(cap(origMessageCh)-noFilteredChanSize, noFilteredChanSize))
+
+	wg.Add(1)
+	go func() { // run messages splitter with filtering
+		defer wg.Done()
+		messagesSplitter(ctx, origMessageCh, noFilteredChan, filteredChan, sm, messageIDFilter)
+	}()
+
+	out := runMessagesMerger(ctx, wg, noFilteredChan, filteredChan, sm)
+
+	lp := aggregatedLenProvider{
+		chanLenProvider[peer.ProtoMessage](origMessageCh),
+		chanLenProvider[protoMessageWrapper](noFilteredChan),
+		chanLenProvider[protoMessageWrapper](filteredChan),
+		chanLenProvider[peer.ProtoMessage](out),
+	}
+	return out, lp, wg
+}
+
+func deduplicateProtoTxMessages(
+	ctx context.Context, origMessageCh <-chan peer.ProtoMessage,
+) (<-chan peer.ProtoMessage, lenProvider, waiter) {
+	payloadIDFilter := txMessagePayloadIDFilter{seed: maphash.MakeSeed()}
+	return deduplicateProtoMessages(ctx, origMessageCh, payloadIDFilter.Filter)
+}

pkg/node/helpers_internal_test.go

@@ -0,0 +1,93 @@
+package node
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/wavesplatform/gowaves/pkg/p2p/peer"
+	"github.com/wavesplatform/gowaves/pkg/proto"
+)
+
+func testContext(t *testing.T) (context.Context, context.CancelFunc) {
+	t.Helper()
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	return ctx, cancel
+}
+
+func filledMsgChan(values ...proto.Message) chan peer.ProtoMessage {
+	ch := make(chan peer.ProtoMessage, len(values))
+	for _, v := range values {
+		ch <- peer.ProtoMessage{
+			ID:      nil, // ID is not used in this test.
+			Message: v,
+		}
+	}
+	return ch
+}
+
+func readProtoMessages(t *testing.T, ch <-chan peer.ProtoMessage, expectedMsgCount int) []proto.Message {
+	timeout := time.NewTimer(5 * time.Second)
+	messages := make([]proto.Message, 0, expectedMsgCount)
+	for range expectedMsgCount {
+		select {
+		case <-timeout.C:
+			t.Fatalf("timed out waiting for messages, check expectedMsgCount=%d", expectedMsgCount)
+		case msg := <-ch:
+			messages = append(messages, msg.Message)
+		}
+	}
+	select {
+	case msg := <-ch:
+		t.Fatalf("unexpected message=%+v received, expected exactly %d messages", msg, expectedMsgCount)
+	default:
+		// No more messages expected, this is fine.
+	}
+	return messages
+}
+
+func TestDeduplicateProtoTxMessages(t *testing.T) {
+	messages := filledMsgChan(
+		&proto.TransactionMessage{Transaction: nil},
+		&proto.TransactionMessage{Transaction: nil},   // duplicate
+		&proto.PBTransactionMessage{Transaction: nil}, // considered as duplicate because payloads are equal
+		&proto.TransactionMessage{Transaction: []byte{1, 2, 3}},
+		&proto.TransactionMessage{Transaction: []byte{1, 2, 3}},   // duplicate
+		&proto.PBTransactionMessage{Transaction: []byte{1, 2, 3}}, // duplicate, payloads are equal
+		&proto.PBTransactionMessage{Transaction: []byte{42}},
+		&proto.TransactionMessage{Transaction: []byte{42}}, // duplicate, payloads are equal
+		// -- non-transaction messages
+		&proto.ScoreMessage{Score: nil},
+		&proto.ScoreMessage{Score: nil},
+		&proto.ScoreMessage{Score: []byte{1, 2, 3}},
+		&proto.ScoreMessage{Score: []byte{1, 2, 3}},
+		&proto.ScoreMessage{Score: []byte{42}},
+		&proto.ScoreMessage{Score: []byte{42}},
+		&proto.ScoreMessage{Score: []byte{21}},
+	)
+
+	expected := []proto.Message{
+		&proto.TransactionMessage{Transaction: nil},
+		&proto.TransactionMessage{Transaction: []byte{1, 2, 3}},
+		&proto.PBTransactionMessage{Transaction: []byte{42}},
+		// -- non-transaction messages
+		&proto.ScoreMessage{Score: nil},
+		&proto.ScoreMessage{Score: nil},
+		&proto.ScoreMessage{Score: []byte{1, 2, 3}},
+		&proto.ScoreMessage{Score: []byte{1, 2, 3}},
+		&proto.ScoreMessage{Score: []byte{42}},
+		&proto.ScoreMessage{Score: []byte{42}},
+		&proto.ScoreMessage{Score: []byte{21}},
+	}
+
+	ctx, cancel := testContext(t)
+	deduplicated, _, wg := deduplicateProtoTxMessages(ctx, messages)
+	defer wg.Wait()
+	defer cancel()
+
+	actual := readProtoMessages(t, deduplicated, len(expected))
+	assert.ElementsMatch(t, expected, actual)
+}

pkg/node/node.go

@@ -156,10 +156,11 @@ func (a *Node) Run(
 	ctx context.Context, p peer.Parent, internalMessageCh <-chan messages.InternalMessage,
 	networkMsgCh <-chan network.InfoMessage, syncPeer *network.SyncPeer,
 ) {
-	protoMessagesChan := chanLenProvider[peer.ProtoMessage](p.MessageCh)
+	messageCh, protoMessagesLenProvider, wg := deduplicateProtoTxMessages(ctx, p.MessageCh)
+	defer wg.Wait()
 
 	go a.runOutgoingConnections(ctx)
-	go a.runInternalMetrics(ctx, protoMessagesChan)
+	go a.runInternalMetrics(ctx, protoMessagesLenProvider)
 	go a.runIncomingConnections(ctx)
 
 	tasksCh := make(chan tasks.AsyncTask, 10)
@@ -207,7 +208,7 @@ func (a *Node) Run(
 			default:
 				zap.S().Warnf("[%s] Unknown network info message '%T'", m.State.State, msg)
 			}
-		case mess := <-p.MessageCh:
+		case mess := <-messageCh:
 			zap.S().Named(logging.FSMNamespace).Debugf("[%s] Network message '%T' received from '%s'",
 				m.State.State, mess.Message, mess.ID.ID())
 			action, ok := actions[reflect.TypeOf(mess.Message)]
@@ -235,10 +236,6 @@ type lenProvider interface {
 	Len() int
 }
 
-type chanLenProvider[T any] <-chan T
-
-func (c chanLenProvider[T]) Len() int { return len(c) }
-
 func (a *Node) runInternalMetrics(ctx context.Context, protoMessagesChan lenProvider) {
 	ticker := time.NewTicker(metricInternalChannelSizeUpdateInterval)
 	defer ticker.Stop()

pkg/p2p/peer/handle.go

@@ -64,7 +64,7 @@ func Handle(ctx context.Context, peer Peer, parent Parent, remote Remote) error
 	var errSentToParent bool // if errSentToParent is true then we need to wait ctx cancellation
 	for {
 		select {
-		case <-ctx.Done():
+		case <-ctx.Done(): // context is unique for each peer, so when passed 'peer' arg is closed, context is canceled
 			//TODO: On Done() Err() contains only Canceled or DeadlineExceeded.
 			// Actually, those errors are only logged in different places and not used to alter behavior.
 			// Consider removing wrapping. For now, if context was canceled no error is passed by.