diff --git a/alpaca/stream.go b/alpaca/stream.go index ac832d3..ac73fc6 100644 --- a/alpaca/stream.go +++ b/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).")) } diff --git a/polygon/stream.go b/polygon/stream.go index a34910c..04e8a37 100644 --- a/polygon/stream.go +++ b/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).")) }