Use context aware routines

Author: mpetrun5

tss/ecdsa/common/base.go

@@ -14,6 +14,7 @@ import (
 	"github.com/libp2p/go-libp2p/core/host"
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/rs/zerolog"
+	"github.com/sourcegraph/conc/pool"
 	"github.com/sprintertech/sprinter-signing/comm"
 	"github.com/sprintertech/sprinter-signing/tss/message"
 )
@@ -57,15 +58,17 @@ func (b *BaseTss) ProcessInboundMessages(ctx context.Context, msgChan chan *comm
 
 	for {
 		select {
-		case wMsg := <-msgChan:
+		case msg := <-msgChan:
 			{
-				go func(wMsg *comm.WrappedMessage) {
+				wMsg := msg
+				p := pool.New().WithContext(ctx).WithCancelOnError()
+				p.Go(func(ctx context.Context) error {
 					b.Log.Debug().Msgf("Processed inbound message from %s", wMsg.From)
 
 					msg, err := message.UnmarshalTssMessage(wMsg.Payload)
 					if err != nil {
 						b.Log.Error().Err(err).Msgf("Failed unmarshaling message from %s", wMsg.From)
-						return
+						return err
 					}
 
 					ok, err := b.Party.UpdateFromBytes(
@@ -75,10 +78,11 @@ func (b *BaseTss) ProcessInboundMessages(ctx context.Context, msgChan chan *comm
 						new(big.Int).SetBytes([]byte(b.SID)))
 					if !ok {
 						b.Log.Error().Err(err).Msgf("Failed updating party with message from %s", wMsg.From)
-						return
+						return err
 					}
 					b.Log.Debug().Msgf("Updated party with message from %s", wMsg.From)
-				}(wMsg)
+					return nil
+				})
 			}
 		case <-ctx.Done():
 			return nil
@@ -93,33 +97,37 @@ func (b *BaseTss) ProcessOutboundMessages(ctx context.Context, outChn chan tss.M
 		select {
 		case msg := <-outChn:
 			{
-				go func(msg tss.Message) {
-					b.Log.Debug().Msg(msg.String())
-					wireBytes, routing, err := msg.WireBytes()
+				wMsg := msg
+				p := pool.New().WithContext(ctx).WithCancelOnError()
+				p.Go(func(ctx context.Context) error {
+					b.Log.Debug().Msg(wMsg.String())
+					wireBytes, routing, err := wMsg.WireBytes()
 					if err != nil {
 						b.Log.Error().Err(err).Msgf("Failed getting wire bytes")
-						return
+						return err
 					}
 
 					msgBytes, err := message.MarshalTssMessage(wireBytes, routing.IsBroadcast)
 					if err != nil {
 						b.Log.Error().Err(err).Msgf("Failed marshaling message")
-						return
+						return err
 					}
 
-					peers, err := b.BroadcastPeers(msg)
+					peers, err := b.BroadcastPeers(wMsg)
 					if err != nil {
 						b.Log.Error().Err(err).Msgf("Failed getting broadcast peers")
-						return
+						return err
 					}
 
 					b.Log.Debug().Msgf("Sending message to %s", peers)
 					err = b.Communication.Broadcast(peers, msgBytes, messageType, b.SessionID())
 					if err != nil {
 						b.Log.Error().Err(err).Msgf("Failed broadcasting message")
-						return
+						return err
 					}
-				}(msg)
+
+					return nil
+				})
 			}
 		case <-ctx.Done():
 			{