From dd0571229bbadba4486af2df6a55d48c1590800b Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Tue, 15 Nov 2022 07:53:50 -0300
Subject: [PATCH] RelayConnect() ensures there will be a connection, and handle
 connection errors better.

---
 relay.go     | 20 +++++++++++++-------
 relaypool.go |  6 +++++-
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/relay.go b/relay.go
index b8932fc..976170e 100644
--- a/relay.go
+++ b/relay.go
@@ -39,15 +39,16 @@ type Relay struct {
 	Connection    *Connection
 	subscriptions s.MapOf[string, *Subscription]
 
-	Notices     chan string
+	Notices         chan string
+	ConnectionError chan error
+
 	statusChans s.MapOf[string, chan Status]
 }
 
-func NewRelay(url string) *Relay {
-	return &Relay{
-		URL:           NormalizeURL(url),
-		subscriptions: s.MapOf[string, *Subscription]{},
-	}
+func RelayConnect(url string) (*Relay, error) {
+	r := &Relay{URL: NormalizeURL(url)}
+	err := r.Connect()
+	return r, err
 }
 
 func (r *Relay) Connect() error {
@@ -60,6 +61,9 @@ func (r *Relay) Connect() error {
 		return fmt.Errorf("error opening websocket to '%s': %w", r.URL, err)
 	}
 
+	r.Notices = make(chan string)
+	r.ConnectionError = make(chan error)
+
 	conn := NewConnection(socket)
 	r.Connection = conn
 
@@ -67,8 +71,10 @@ func (r *Relay) Connect() error {
 		for {
 			typ, message, err := conn.socket.ReadMessage()
 			if err != nil {
-				continue
+				r.ConnectionError <- err
+				break
 			}
+
 			if typ == websocket.PingMessage {
 				conn.WriteMessage(websocket.PongMessage, nil)
 				continue
diff --git a/relaypool.go b/relaypool.go
index b615dc5..4c76ecf 100644
--- a/relaypool.go
+++ b/relaypool.go
@@ -65,7 +65,11 @@ func (r *RelayPool) Add(url string, policy RelayPoolPolicy) error {
 		policy = SimplePolicy{Read: true, Write: true}
 	}
 
-	relay := NewRelay(url)
+	relay, err := RelayConnect(url)
+	if err != nil {
+		return err
+	}
+
 	r.Policies.Store(relay.URL, policy)
 	r.Relays.Store(relay.URL, relay)