@@ -76,12 +76,6 @@ func (s *Server) Start() {
7676				"connection_id" , s .callbackContext .ConnectionId .String (),
7777			)
7878		s .Protocol .Start ()
79- 		// Start goroutine to cleanup resources on protocol shutdown 
80- 		go  func () {
81- 			<- s .Protocol .DoneChan ()
82- 			close (s .requestTxIdsResultChan )
83- 			close (s .requestTxsResultChan )
84- 		}()
8579	})
8680}
8781
@@ -103,13 +97,14 @@ func (s *Server) RequestTxIds(
10397		return  nil , err 
10498	}
10599	// Wait for result 
106- 	txIds ,  ok   :=   <- s . requestTxIdsResultChan 
107- 	if   ! ok  { 
100+ 	select  { 
101+ 	case   <- s . DoneChan (): 
108102		return  nil , protocol .ProtocolShuttingDownError 
103+ 	case  txIds  :=  <- s .requestTxIdsResultChan :
104+ 		// Update ack count for next call 
105+ 		s .ackCount  =  len (txIds )
106+ 		return  txIds , nil 
109107	}
110- 	// Update ack count for next call 
111- 	s .ackCount  =  len (txIds )
112- 	return  txIds , nil 
113108}
114109
115110// RequestTxs requests the content of the requested TX identifiers from the remote node's mempool 
@@ -127,11 +122,12 @@ func (s *Server) RequestTxs(txIds []TxId) ([]TxBody, error) {
127122		return  nil , err 
128123	}
129124	// Wait for result 
130- 	txs ,  ok   :=   <- s . requestTxsResultChan 
131- 	if   ! ok  { 
125+ 	select  { 
126+ 	case   <- s . DoneChan (): 
132127		return  nil , protocol .ProtocolShuttingDownError 
128+ 	case  txs  :=  <- s .requestTxsResultChan :
129+ 		return  txs , nil 
133130	}
134- 	return  txs , nil 
135131}
136132
137133func  (s  * Server ) messageHandler (msg  protocol.Message ) error  {
@@ -163,12 +159,6 @@ func (s *Server) handleReplyTxIds(msg protocol.Message) error {
163159			"role" , "server" ,
164160			"connection_id" , s .callbackContext .ConnectionId .String (),
165161		)
166- 	// Check for shutdown 
167- 	select  {
168- 	case  <- s .Protocol .DoneChan ():
169- 		return  protocol .ProtocolShuttingDownError 
170- 	default :
171- 	}
172162	msgReplyTxIds  :=  msg .(* MsgReplyTxIds )
173163	s .requestTxIdsResultChan  <-  msgReplyTxIds .TxIds 
174164	return  nil 
@@ -182,12 +172,6 @@ func (s *Server) handleReplyTxs(msg protocol.Message) error {
182172			"role" , "server" ,
183173			"connection_id" , s .callbackContext .ConnectionId .String (),
184174		)
185- 	// Check for shutdown 
186- 	select  {
187- 	case  <- s .Protocol .DoneChan ():
188- 		return  protocol .ProtocolShuttingDownError 
189- 	default :
190- 	}
191175	msgReplyTxs  :=  msg .(* MsgReplyTxs )
192176	s .requestTxsResultChan  <-  msgReplyTxs .Txs 
193177	return  nil 
0 commit comments