From c0a42ff2ff663f75ba3115e7bc76d3e37092e2c6 Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Sat, 6 May 2023 14:32:39 -0300
Subject: [PATCH] SimplePool improvements.

---
 pool.go | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/pool.go b/pool.go
index 8de7142..1007a43 100644
--- a/pool.go
+++ b/pool.go
@@ -8,14 +8,21 @@ import (
 )
 
 type SimplePool struct {
-	Relays map[string]*Relay
+	Relays  map[string]*Relay
+	Context context.Context
 
-	mutex sync.Mutex
+	mutex  sync.Mutex
+	cancel context.CancelFunc
 }
 
-func NewSimplePool() *SimplePool {
+func NewSimplePool(ctx context.Context) *SimplePool {
+	ctx, cancel := context.WithCancel(ctx)
+
 	return &SimplePool{
 		Relays: make(map[string]*Relay),
+
+		Context: ctx,
+		cancel:  cancel,
 	}
 }
 
@@ -26,13 +33,13 @@ func (pool *SimplePool) EnsureRelay(url string) *Relay {
 	defer pool.mutex.Unlock()
 
 	relay, ok := pool.Relays[nm]
-	if ok {
+	if ok && relay.ConnectionContext.Err() == nil {
 		// already connected, unlock and return
 		return relay
 	} else {
 		var err error
-		// when connecting to a relay we want the connection to persist forever if possible, so use a new context
-		relay, err = RelayConnect(context.Background(), nm)
+		// we use this ctx here so when the pool dies everything dies
+		relay, err = RelayConnect(pool.Context, nm)
 		if err != nil {
 			return nil
 		}