Merge pull request #383 from sonicfromnewyoke/fix/ws-race-and-leak-fixes

fix: ws race and leak fixes

Author: HealthyBuilder

rpc/sendAndConfirmTransaction/sendAndConfirmTransaction.go

@@ -120,25 +120,18 @@ func WaitForConfirmation(
 		timeout = &t
 	}
 
-	for {
-		select {
-		case <-ctx.Done():
-			return false, ctx.Err()
-		case <-time.After(*timeout):
+	timeoutCtx, cancel := context.WithTimeout(ctx, *timeout)
+	defer cancel()
+
+	got, err := sub.Recv(timeoutCtx)
+	if err != nil {
+		if timeoutCtx.Err() == context.DeadlineExceeded {
 			return false, ErrTimeout
-		case resp, ok := <-sub.Response():
-			if !ok {
-				return false, fmt.Errorf("subscription closed")
-			}
-			if resp.Value.Err != nil {
-				// The transaction was confirmed, but it failed while executing (one of the instructions failed).
-				return true, fmt.Errorf("confirmed transaction with execution error: %v", resp.Value.Err)
-			} else {
-				// Success! Confirmed! And there was no error while executing the transaction.
-				return true, nil
-			}
-		case err := <-sub.Err():
-			return false, err
 		}
+		return false, err
+	}
+	if got.Value.Err != nil {
+		return true, fmt.Errorf("confirmed transaction with execution error: %v", got.Value.Err)
 	}
+	return true, nil
 }

rpc/ws/accountSubscribe.go

@@ -101,19 +101,6 @@ func (sw *AccountSubscription) Err() <-chan error {
 	return sw.sub.err
 }
 
-func (sw *AccountSubscription) Response() <-chan *AccountResult {
-	typedChan := make(chan *AccountResult, 1)
-	go func(ch chan *AccountResult) {
-		// TODO: will this subscription yield more than one result?
-		d, ok := <-sw.sub.stream
-		if !ok {
-			return
-		}
-		ch <- d.(*AccountResult)
-	}(typedChan)
-	return typedChan
-}
-
 func (sw *AccountSubscription) Unsubscribe() {
 	sw.sub.Unsubscribe()
 }

rpc/ws/blockSubscribe.go

@@ -167,19 +167,6 @@ func (sw *BlockSubscription) Err() <-chan error {
 	return sw.sub.err
 }
 
-func (sw *BlockSubscription) Response() <-chan *BlockResult {
-	typedChan := make(chan *BlockResult, 1)
-	go func(ch chan *BlockResult) {
-		// TODO: will this subscription yield more than one result?
-		d, ok := <-sw.sub.stream
-		if !ok {
-			return
-		}
-		ch <- d.(*BlockResult)
-	}(typedChan)
-	return typedChan
-}
-
 func (sw *BlockSubscription) Unsubscribe() {
 	sw.sub.Unsubscribe()
 }

rpc/ws/client.go

@@ -42,10 +42,10 @@ type Client struct {
 	conn                    *websocket.Conn
 	connCtx                 context.Context
 	connCtxCancel           context.CancelFunc
+	wg                      sync.WaitGroup
 	lock                    sync.RWMutex
 	subscriptionByRequestID map[uint64]*Subscription
 	subscriptionByWSSubID   map[uint64]*Subscription
-	reconnectOnErr          bool
 	shortID                 bool
 }
 
@@ -105,10 +105,13 @@ func ConnectWithOptions(ctx context.Context, rpcEndpoint string, opt *Options) (
 	}
 
 	c.connCtx, c.connCtxCancel = context.WithCancel(context.Background())
+	c.wg.Add(2)
 	go func() {
+		defer c.wg.Done()
 		c.conn.SetReadDeadline(time.Now().Add(pongWait))
 		c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
 		ticker := time.NewTicker(pingPeriod)
+		defer ticker.Stop()
 		for {
 			select {
 			case <-c.connCtx.Done():
@@ -132,22 +135,27 @@ func (c *Client) sendPing() {
 	}
 }
 
+// Close cancels the connection context, closes the underlying websocket
+// connection, and waits for background goroutines to finish.
 func (c *Client) Close() {
 	c.lock.Lock()
-	defer c.lock.Unlock()
 	c.connCtxCancel()
 	c.conn.Close()
+	c.lock.Unlock()
+	c.wg.Wait()
 }
 
 func (c *Client) receiveMessages() {
+	defer c.wg.Done()
+	defer c.closeAllSubscription(ErrSubscriptionClosed)
+
 	for {
 		select {
 		case <-c.connCtx.Done():
 			return
 		default:
 			_, message, err := c.conn.ReadMessage()
 			if err != nil {
-				c.closeAllSubscription(err)
 				return
 			}
 			c.handleMessage(message)
@@ -218,7 +226,6 @@ func (c *Client) handleNewSubscriptionMessage(requestID, subID uint64) {
 		zap.Uint64("request_id", requestID),
 		zap.Int("subscription_count", len(c.subscriptionByWSSubID)),
 	)
-	return
 }
 
 func (c *Client) handleSubscriptionMessage(subID uint64, message []byte) {
@@ -239,34 +246,43 @@ func (c *Client) handleSubscriptionMessage(subID uint64, message []byte) {
 	// Decode the message using the subscription-provided decoderFunc.
 	result, err := sub.decoderFunc(message)
 	if err != nil {
-		fmt.Println("*****************************")
 		c.closeSubscription(sub.req.ID, fmt.Errorf("unable to decode client response: %w", err))
 		return
 	}
 
 	// this cannot be blocking or else
 	// we  will no read any other message
 	if len(sub.stream) >= cap(sub.stream) {
-		zlog.Warn("closing ws client subscription... not consuming fast en ought",
+		zlog.Warn("closing ws client subscription... not consuming fast enough",
 			zap.Uint64("request_id", sub.req.ID),
 		)
 		c.closeSubscription(sub.req.ID, fmt.Errorf("reached channel max capacity %d", len(sub.stream)))
 		return
 	}
 
-	sub.mutex.Lock()
-	defer sub.mutex.Unlock()
+	sub.mu.Lock()
 	if !sub.closed {
 		sub.stream <- result
 	}
+	sub.mu.Unlock()
 }
 
 func (c *Client) closeAllSubscription(err error) {
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
 	for _, sub := range c.subscriptionByRequestID {
-		sub.err <- err
+		sub.mu.Lock()
+		if !sub.closed {
+			select {
+			case sub.err <- err:
+			default:
+			}
+			sub.closed = true
+			close(sub.stream)
+			close(sub.err)
+		}
+		sub.mu.Unlock()
 	}
 
 	c.subscriptionByRequestID = map[uint64]*Subscription{}
@@ -282,12 +298,21 @@ func (c *Client) closeSubscription(reqID uint64, err error) {
 		return
 	}
 
-	sub.err <- err
+	sub.mu.Lock()
+	if !sub.closed {
+		select {
+		case sub.err <- err:
+		default:
+		}
+		sub.closed = true
+		close(sub.stream)
+		close(sub.err)
+	}
+	sub.mu.Unlock()
 
-	err = c.unsubscribe(sub.subID, sub.unsubscribeMethod)
-	if err != nil {
+	if e := c.unsubscribe(sub.subID, sub.unsubscribeMethod); e != nil {
 		zlog.Warn("unable to send rpc unsubscribe call",
-			zap.Error(err),
+			zap.Error(e),
 		)
 	}