net: dispatcher.go. support enable & disable duplicate msg filter for…

… subscribers.
nebulasio · Jan 30, 2018 · 5adbb32 · 5adbb32
1 parent dfed816
commit 5adbb32
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 20 deletions.
diff --git a/cmd/network/main.go b/cmd/network/main.go
@@ -118,8 +118,8 @@ func run(mode, configPath string, packageSize, concurrentMessageCount, totalMess
 	}
 
 	// register dispatcher.
-	netService.Register(net.NewSubscriber(netService, messageCh, PingMessage))
-	netService.Register(net.NewSubscriber(netService, messageCh, PongMessage))
+	netService.Register(net.NewSubscriber(netService, messageCh, false, PingMessage))
+	netService.Register(net.NewSubscriber(netService, messageCh, false, PongMessage))
 
 	// start server.
 	netService.Start()

diff --git a/core/block_pool.go b/core/block_pool.go
@@ -99,9 +99,9 @@ func NewBlockPool(size int) (*BlockPool, error) {
 
 // RegisterInNetwork register message subscriber in network.
 func (pool *BlockPool) RegisterInNetwork(ns net.Service) {
-	ns.Register(net.NewSubscriber(pool, pool.receiveBlockMessageCh, MessageTypeNewBlock))
-	ns.Register(net.NewSubscriber(pool, pool.receiveBlockMessageCh, MessageTypeDownloadedBlockReply))
-	ns.Register(net.NewSubscriber(pool, pool.receiveDownloadBlockMessageCh, MessageTypeDownloadedBlock))
+	ns.Register(net.NewSubscriber(pool, pool.receiveBlockMessageCh, true, MessageTypeNewBlock))
+	ns.Register(net.NewSubscriber(pool, pool.receiveBlockMessageCh, true, MessageTypeDownloadedBlockReply))
+	ns.Register(net.NewSubscriber(pool, pool.receiveDownloadBlockMessageCh, true, MessageTypeDownloadedBlock))
 	pool.ns = ns
 }
 

diff --git a/core/transaction_pool.go b/core/transaction_pool.go
@@ -95,7 +95,7 @@ func (pool *TransactionPool) SetGasConfig(gasPrice, gasLimit *util.Uint128) {
 
 // RegisterInNetwork register message subscriber in network.
 func (pool *TransactionPool) RegisterInNetwork(ns net.Service) {
-	ns.Register(net.NewSubscriber(pool, pool.receivedMessageCh, MessageTypeNewTx))
+	ns.Register(net.NewSubscriber(pool, pool.receivedMessageCh, true, MessageTypeNewTx))
 	pool.ns = ns
 }
 

diff --git a/net/dispatcher.go b/net/dispatcher.go
@@ -40,6 +40,7 @@ type Dispatcher struct {
 	quitCh             chan bool
 	receivedMessageCh  chan Message
 	dispatchedMessages *lru.Cache
+	filters            map[string]bool
 }
 
 // NewDispatcher create Dispatcher instance.
@@ -48,6 +49,7 @@ func NewDispatcher() *Dispatcher {
 		subscribersMap:    new(sync.Map),
 		quitCh:            make(chan bool, 10),
 		receivedMessageCh: make(chan Message, 65536),
+		filters:           make(map[string]bool),
 	}
 
 	dp.dispatchedMessages, _ = lru.New(10240)
@@ -61,6 +63,7 @@ func (dp *Dispatcher) Register(subscribers ...*Subscriber) {
 		for _, mt := range v.msgTypes {
 			m, _ := dp.subscribersMap.LoadOrStore(mt, new(sync.Map))
 			m.(*sync.Map).Store(v, true)
+			dp.filters[mt] = v.DoFilter()
 		}
 	}
 }
@@ -76,6 +79,7 @@ func (dp *Dispatcher) Deregister(subscribers ...*Subscriber) {
 			}
 			m.(*sync.Map).Delete(v)
 			dp.subscribersMap.Delete(mt)
+			delete(dp.filters, mt)
 		}
 	}
 }
@@ -134,10 +138,12 @@ func (dp *Dispatcher) Stop() {
 func (dp *Dispatcher) PutMessage(msg Message) {
 	// it's a optimize strategy for message dispatch, according to https://github.com/nebulasio/go-nebulas/issues/50
 	hash := msg.Hash()
-	if exist, _ := dp.dispatchedMessages.ContainsOrAdd(hash, hash); exist == true {
-		// duplicated message, ignore.
-		metricsDuplicatedMessage(msg.MessageType())
-		return
+	if dp.filters[msg.MessageType()] {
+		if exist, _ := dp.dispatchedMessages.ContainsOrAdd(hash, hash); exist == true {
+			// duplicated message, ignore.
+			metricsDuplicatedMessage(msg.MessageType())
+			return
+		}
 	}
 
 	dp.receivedMessageCh <- msg

diff --git a/net/types.go b/net/types.go
@@ -106,11 +106,14 @@ type Subscriber struct {
 
 	// msgType message types to subscribe
 	msgTypes []string
+
+	// doFilter dup message
+	doFilter bool
 }
 
 // NewSubscriber return new Subscriber instance.
-func NewSubscriber(id interface{}, msgChan chan Message, msgTypes ...string) *Subscriber {
-	return &Subscriber{id, msgChan, msgTypes}
+func NewSubscriber(id interface{}, msgChan chan Message, doFilter bool, msgTypes ...string) *Subscriber {
+	return &Subscriber{id, msgChan, msgTypes, doFilter}
 }
 
 // ID return id.
@@ -128,6 +131,11 @@ func (s *Subscriber) MessageChan() chan Message {
 	return s.msgChan
 }
 
+// DoFilter return doFilter
+func (s *Subscriber) DoFilter() bool {
+	return s.doFilter
+}
+
 // BaseMessage base message
 type BaseMessage struct {
 	t    string

diff --git a/sync/sync_service.go b/sync/sync_service.go
@@ -69,10 +69,10 @@ func (ss *Service) Start() {
 
 	// register the network handler.
 	netService := ss.netService
-	netService.Register(net.NewSubscriber(ss, ss.messageCh, net.ChainSync))
-	netService.Register(net.NewSubscriber(ss, ss.messageCh, net.ChainChunks))
-	netService.Register(net.NewSubscriber(ss, ss.messageCh, net.ChainGetChunk))
-	netService.Register(net.NewSubscriber(ss, ss.messageCh, net.ChainChunkData))
+	netService.Register(net.NewSubscriber(ss, ss.messageCh, false, net.ChainSync))
+	netService.Register(net.NewSubscriber(ss, ss.messageCh, false, net.ChainChunks))
+	netService.Register(net.NewSubscriber(ss, ss.messageCh, false, net.ChainGetChunk))
+	netService.Register(net.NewSubscriber(ss, ss.messageCh, false, net.ChainChunkData))
 
 	// start loop().
 	go ss.startLoop()
@@ -82,10 +82,10 @@ func (ss *Service) Start() {
 func (ss *Service) Stop() {
 	// deregister the network handler.
 	netService := ss.netService
-	netService.Deregister(net.NewSubscriber(ss, ss.messageCh, net.ChainSync))
-	netService.Deregister(net.NewSubscriber(ss, ss.messageCh, net.ChainChunks))
-	netService.Deregister(net.NewSubscriber(ss, ss.messageCh, net.ChainGetChunk))
-	netService.Deregister(net.NewSubscriber(ss, ss.messageCh, net.ChainChunkData))
+	netService.Deregister(net.NewSubscriber(ss, ss.messageCh, false, net.ChainSync))
+	netService.Deregister(net.NewSubscriber(ss, ss.messageCh, false, net.ChainChunks))
+	netService.Deregister(net.NewSubscriber(ss, ss.messageCh, false, net.ChainGetChunk))
+	netService.Deregister(net.NewSubscriber(ss, ss.messageCh, false, net.ChainChunkData))
 
 	ss.StopActiveSync()