expose output channel and guard against user closing it (#103)

Author: morningvera

websocket/polygon.go

@@ -154,17 +154,10 @@ func (c *Client) Unsubscribe(topic Topic, tickers ...string) error {
 	return nil
 }
 
-// Output returns the next message in the output queue. If no messages are available, it returns nil.
-func (c *Client) Output() any {
-	select {
-	case out, ok := <-c.output:
-		if ok {
-			return out
-		}
-		return nil
-	default:
-		return nil
-	}
+// Output returns the output queue. If the channel is closed by the user (not
+// recommended), the client connection will close as well.
+func (c *Client) Output() chan any {
+	return c.output
 }
 
 // Close attempt to gracefully close the connection to the server.
@@ -254,6 +247,17 @@ func (c *Client) reconnect() {
 	}
 }
 
+func (c *Client) closeOutput() {
+	defer func() {
+		if r := recover(); r != nil {
+			c.log.Debugf("output channel was closed by user")
+		} else {
+			c.log.Debugf("output channel closed")
+		}
+	}()
+	close(c.output)
+}
+
 func (c *Client) close(reconnect bool) {
 	if c.conn == nil {
 		return
@@ -270,7 +274,7 @@ func (c *Client) close(reconnect bool) {
 			c.log.Errorf("process thread closed: %v", err)
 		}
 		c.shouldClose = true
-		close(c.output)
+		c.closeOutput()
 	}
 
 	if c.conn != nil {
@@ -329,16 +333,15 @@ func (c *Client) write() error {
 			if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
 				return fmt.Errorf("failed to send message: %w", err)
 			}
-		default:
-			continue
 		}
 	}
 }
 
 func (c *Client) process() (err error) {
 	defer func() {
 		c.log.Debugf("process thread closed")
-		if err != nil {
+		r := recover()
+		if r != nil || err != nil {
 			go c.Close() // this client should close if it hits a fatal error (e.g. auth failed)
 		}
 	}()
@@ -356,8 +359,6 @@ func (c *Client) process() (err error) {
 			if err := c.route(msgs); err != nil {
 				return err
 			}
-		default:
-			continue
 		}
 	}
 }

websocket/polygon_live_test.go

@@ -51,14 +51,16 @@ func TestMain(t *testing.T) {
 		log.Fatal(err)
 	}
 
-	go printOutput(ctx, c) // comment for raw data handling
-	// go printRawOutput(ctx, c) // uncomment for raw data handling
+	go printOutput(c) // comment for raw data handling
+	// go printRawOutput(c) // uncomment for raw data handling
 
 	time.Sleep(10 * time.Second)
 	if err := c.Subscribe(polygonws.StocksTrades, "*"); err != nil {
 		log.Error(err)
 	}
 
+	// close(c.Output()) // uncomment to test behavior when output channel is closed
+
 	time.Sleep(250 * time.Millisecond)
 	if err := c.Unsubscribe(polygonws.StocksTrades); err != nil {
 		log.Error(err)
@@ -91,35 +93,17 @@ func TestMain(t *testing.T) {
 	}
 }
 
-func printOutput(ctx context.Context, client *polygonws.Client) {
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		default:
-			out := client.Output()
-			if out == nil {
-				continue
-			}
-			fmt.Println(out)
-		}
+func printOutput(client *polygonws.Client) {
+	for out := range client.Output() {
+		fmt.Println(out)
 	}
 }
 
 //nolint:deadcode
-func printRawOutput(ctx context.Context, client *polygonws.Client) {
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		default:
-			out := client.Output()
-			if out == nil {
-				continue
-			}
-			if b, ok := out.(json.RawMessage); ok {
-				fmt.Println(string(b))
-			}
+func printRawOutput(client *polygonws.Client) {
+	for out := range client.Output() {
+		if b, ok := out.(json.RawMessage); ok {
+			fmt.Println(string(b))
 		}
 	}
 }

websocket/polygon_test.go

@@ -89,10 +89,6 @@ func TestConnectAuthSuccess(t *testing.T) {
 	// closing before connecting shouldn't do anthing
 	c.Close()
 
-	// accessing output early shouldn't do anything
-	out := c.Output()
-	assert.Nil(t, out)
-
 	// connect successfully
 	err = c.Connect()
 	assert.Nil(t, err)