Merge pull request #52 from alpacahq/feature/websocket-connect-failure

Add delayed reconnects on socket open failures

alpaca/stream.go

@@ -19,6 +19,10 @@ const (
 	AccountUpdates = "account_updates"
 )
 
+const (
+	MaxConnectionAttempts = 3
+)
+
 var (
 	once sync.Once
 	str  *Stream
@@ -217,9 +221,17 @@ func openSocket() *websocket.Conn {
 		scheme = "ws"
 	}
 	u := url.URL{Scheme: scheme, Host: ub.Host, Path: "/stream"}
-	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
-	if err != nil {
-		panic(err)
+	connectionAttempts := 0
+	for connectionAttempts < MaxConnectionAttempts {
+		connectionAttempts++
+		c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+		if err == nil {
+			return c
+		}
+		if connectionAttempts == MaxConnectionAttempts {
+			panic(err)
+		}
+		time.Sleep(1 * time.Second)
 	}
-	return c
+	panic(fmt.Errorf("Error: Could not open Alpaca stream (max retries exceeded)."))
 }

polygon/stream.go

@@ -21,6 +21,10 @@ const (
 	Quotes     = "Q"
 )
 
+const (
+	MaxConnectionAttempts = 3
+)
+
 var (
 	once sync.Once
 	str  *Stream
@@ -238,14 +242,22 @@ func openSocket() *websocket.Conn {
 	if !ok {
 		polygonStreamEndpoint = "wss://alpaca.socket.polygon.io/stocks"
 	}
-	c, _, err := websocket.DefaultDialer.Dial(polygonStreamEndpoint, nil)
-	if err != nil {
-		panic(err)
-	}
-	// read connection message
-	msg := []PolgyonServerMsg{}
-	if err = c.ReadJSON(&msg); err != nil {
-		panic(err)
+	connectionAttempts := 0
+	for connectionAttempts < MaxConnectionAttempts {
+		connectionAttempts++
+		c, _, err := websocket.DefaultDialer.Dial(polygonStreamEndpoint, nil)
+		if err != nil {
+			if connectionAttempts == MaxConnectionAttempts {
+				panic(err)
+			}
+		} else {
+			// consume connection message
+			msg := []PolgyonServerMsg{}
+			if err = c.ReadJSON(&msg); err != nil {
+				return c
+			}
+		}
+		time.Sleep(1 * time.Second)
 	}
-	return c
+	panic(fmt.Errorf("Error: Could not open Polygon stream (max retries exceeded)."))
 }