From fc4b37f7bcdc2de36f279c458ce79da312d8d29e Mon Sep 17 00:00:00 2001
From: Zoltan Papp <zoltan.pmail@gmail.com>
Date: Thu, 19 Sep 2024 13:49:28 +0200
Subject: [PATCH] Exit from processConnResults after all tries (#2621)

* Exit from processConnResults after all tries

If all server is unavailable then the server picker never return
because we never close the result channel.
Count the number of the results and exit when we reached the
expected size
---
 relay/client/picker.go      | 10 +++++++---
 relay/client/picker_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)
 create mode 100644 relay/client/picker_test.go

diff --git a/relay/client/picker.go b/relay/client/picker.go
index b0888a4a0..13b0547aa 100644
--- a/relay/client/picker.go
+++ b/relay/client/picker.go
@@ -35,12 +35,15 @@ func (sp *ServerPicker) PickServer(parentCtx context.Context, urls []string) (*C
 
 	connResultChan := make(chan connResult, totalServers)
 	successChan := make(chan connResult, 1)
-
 	concurrentLimiter := make(chan struct{}, maxConcurrentServers)
+
 	for _, url := range urls {
+		// todo check if we have a successful connection so we do not need to connect to other servers
 		concurrentLimiter <- struct{}{}
 		go func(url string) {
-			defer func() { <-concurrentLimiter }()
+			defer func() {
+				<-concurrentLimiter
+			}()
 			sp.startConnection(parentCtx, connResultChan, url)
 		}(url)
 	}
@@ -72,7 +75,8 @@ func (sp *ServerPicker) startConnection(ctx context.Context, resultChan chan con
 
 func (sp *ServerPicker) processConnResults(resultChan chan connResult, successChan chan connResult) {
 	var hasSuccess bool
-	for cr := range resultChan {
+	for numOfResults := 0; numOfResults < cap(resultChan); numOfResults++ {
+		cr := <-resultChan
 		if cr.Err != nil {
 			log.Debugf("failed to connect to Relay server: %s: %v", cr.Url, cr.Err)
 			continue
diff --git a/relay/client/picker_test.go b/relay/client/picker_test.go
new file mode 100644
index 000000000..f5649d700
--- /dev/null
+++ b/relay/client/picker_test.go
@@ -0,0 +1,31 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+)
+
+func TestServerPicker_UnavailableServers(t *testing.T) {
+	sp := ServerPicker{
+		TokenStore: nil,
+		PeerID:     "test",
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+
+	go func() {
+		_, err := sp.PickServer(ctx, []string{"rel://dummy1", "rel://dummy2"})
+		if err == nil {
+			t.Error(err)
+		}
+		cancel()
+	}()
+
+	<-ctx.Done()
+	if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+		t.Errorf("PickServer() took too long to complete")
+	}
+}