Refactor healthcheck sender and receiver to use configurable options (#4433)

Author: pappz

.github/workflows/golang-test-linux.yml

@@ -217,7 +217,7 @@ jobs:
           - arch: "386"
             raceFlag: ""
           - arch: "amd64"
-            raceFlag: ""
+            raceFlag: "-race"
     runs-on: ubuntu-22.04
     steps:
       - name: Install Go

shared/relay/client/manager.go

@@ -78,9 +78,10 @@ func NewManager(ctx context.Context, serverURLs []string, peerID string, mtu uin
 		tokenStore: tokenStore,
 		mtu:        mtu,
 		serverPicker: &ServerPicker{
-			TokenStore: tokenStore,
-			PeerID:     peerID,
-			MTU:        mtu,
+			TokenStore:        tokenStore,
+			PeerID:            peerID,
+			MTU:               mtu,
+			ConnectionTimeout: defaultConnectionTimeout,
 		},
 		relayClients:            make(map[string]*RelayTrack),
 		onDisconnectedListeners: make(map[string]*list.List),

shared/relay/client/picker.go

@@ -13,11 +13,8 @@ import (
 )
 
 const (
-	maxConcurrentServers = 7
-)
-
-var (
-	connectionTimeout = 30 * time.Second
+	maxConcurrentServers     = 7
+	defaultConnectionTimeout = 30 * time.Second
 )
 
 type connResult struct {
@@ -27,14 +24,15 @@ type connResult struct {
 }
 
 type ServerPicker struct {
-	TokenStore *auth.TokenStore
-	ServerURLs atomic.Value
-	PeerID     string
-	MTU        uint16
+	TokenStore        *auth.TokenStore
+	ServerURLs        atomic.Value
+	PeerID            string
+	MTU               uint16
+	ConnectionTimeout time.Duration
 }
 
 func (sp *ServerPicker) PickServer(parentCtx context.Context) (*Client, error) {
-	ctx, cancel := context.WithTimeout(parentCtx, connectionTimeout)
+	ctx, cancel := context.WithTimeout(parentCtx, sp.ConnectionTimeout)
 	defer cancel()
 
 	totalServers := len(sp.ServerURLs.Load().([]string))

shared/relay/client/picker_test.go

@@ -8,15 +8,15 @@ import (
 )
 
 func TestServerPicker_UnavailableServers(t *testing.T) {
-	connectionTimeout = 5 * time.Second
-
+	timeout := 5 * time.Second
 	sp := ServerPicker{
-		TokenStore: nil,
-		PeerID:     "test",
+		TokenStore:        nil,
+		PeerID:            "test",
+		ConnectionTimeout: timeout,
 	}
 	sp.ServerURLs.Store([]string{"rel://dummy1", "rel://dummy2"})
 
-	ctx, cancel := context.WithTimeout(context.Background(), connectionTimeout+1)
+	ctx, cancel := context.WithTimeout(context.Background(), timeout+1)
 	defer cancel()
 
 	go func() {

shared/relay/healthcheck/env.go

@@ -0,0 +1,24 @@
+package healthcheck
+
+import (
+	"os"
+	"strconv"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	defaultAttemptThresholdEnv = "NB_RELAY_HC_ATTEMPT_THRESHOLD"
+)
+
+func getAttemptThresholdFromEnv() int {
+	if attemptThreshold := os.Getenv(defaultAttemptThresholdEnv); attemptThreshold != "" {
+		threshold, err := strconv.ParseInt(attemptThreshold, 10, 64)
+		if err != nil {
+			log.Errorf("Failed to parse attempt threshold from environment variable \"%s\" should be an integer. Using default value", attemptThreshold)
+			return defaultAttemptThreshold
+		}
+		return int(threshold)
+	}
+	return defaultAttemptThreshold
+}

shared/relay/healthcheck/env_test.go

@@ -0,0 +1,36 @@
+package healthcheck
+
+import (
+	"os"
+	"testing"
+)
+
+//nolint:tenv
+func TestGetAttemptThresholdFromEnv(t *testing.T) {
+	tests := []struct {
+		name     string
+		envValue string
+		expected int
+	}{
+		{"Default attempt threshold when env is not set", "", defaultAttemptThreshold},
+		{"Custom attempt threshold when env is set to a valid integer", "3", 3},
+		{"Default attempt threshold when env is set to an invalid value", "invalid", defaultAttemptThreshold},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.envValue == "" {
+				os.Unsetenv(defaultAttemptThresholdEnv)
+			} else {
+				os.Setenv(defaultAttemptThresholdEnv, tt.envValue)
+			}
+
+			result := getAttemptThresholdFromEnv()
+			if result != tt.expected {
+				t.Fatalf("Expected %d, got %d", tt.expected, result)
+			}
+
+			os.Unsetenv(defaultAttemptThresholdEnv)
+		})
+	}
+}

shared/relay/healthcheck/receiver.go

@@ -7,10 +7,15 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-var (
-	heartbeatTimeout = healthCheckInterval + 10*time.Second
+const (
+	defaultHeartbeatTimeout = defaultHealthCheckInterval + 10*time.Second
 )
 
+type ReceiverOptions struct {
+	HeartbeatTimeout time.Duration
+	AttemptThreshold int
+}
+
 // Receiver is a healthcheck receiver
 // It will listen for heartbeat and check if the heartbeat is not received in a certain time
 // If the heartbeat is not received in a certain time, it will send a timeout signal and stop to work
@@ -27,6 +32,23 @@ type Receiver struct {
 
 // NewReceiver creates a new healthcheck receiver and start the timer in the background
 func NewReceiver(log *log.Entry) *Receiver {
+	opts := ReceiverOptions{
+		HeartbeatTimeout: defaultHeartbeatTimeout,
+		AttemptThreshold: getAttemptThresholdFromEnv(),
+	}
+	return NewReceiverWithOpts(log, opts)
+}
+
+func NewReceiverWithOpts(log *log.Entry, opts ReceiverOptions) *Receiver {
+	heartbeatTimeout := opts.HeartbeatTimeout
+	if heartbeatTimeout <= 0 {
+		heartbeatTimeout = defaultHeartbeatTimeout
+	}
+	attemptThreshold := opts.AttemptThreshold
+	if attemptThreshold <= 0 {
+		attemptThreshold = defaultAttemptThreshold
+	}
+
 	ctx, ctxCancel := context.WithCancel(context.Background())
 
 	r := &Receiver{
@@ -35,10 +57,10 @@ func NewReceiver(log *log.Entry) *Receiver {
 		ctx:              ctx,
 		ctxCancel:        ctxCancel,
 		heartbeat:        make(chan struct{}, 1),
-		attemptThreshold: getAttemptThresholdFromEnv(),
+		attemptThreshold: attemptThreshold,
 	}
 
-	go r.waitForHealthcheck()
+	go r.waitForHealthcheck(heartbeatTimeout)
 	return r
 }
 
@@ -55,7 +77,7 @@ func (r *Receiver) Stop() {
 	r.ctxCancel()
 }
 
-func (r *Receiver) waitForHealthcheck() {
+func (r *Receiver) waitForHealthcheck(heartbeatTimeout time.Duration) {
 	ticker := time.NewTicker(heartbeatTimeout)
 	defer ticker.Stop()
 	defer r.ctxCancel()

shared/relay/healthcheck/receiver_test.go

@@ -2,31 +2,18 @@ package healthcheck
 
 import (
 	"context"
-	"fmt"
-	"os"
-	"sync"
 	"testing"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 )
 
-// Mutex to protect global variable access in tests
-var testMutex sync.Mutex
-
 func TestNewReceiver(t *testing.T) {
-	testMutex.Lock()
-	originalTimeout := heartbeatTimeout
-	heartbeatTimeout = 5 * time.Second
-	testMutex.Unlock()
-
-	defer func() {
-		testMutex.Lock()
-		heartbeatTimeout = originalTimeout
-		testMutex.Unlock()
-	}()
-
-	r := NewReceiver(log.WithContext(context.Background()))
+
+	opts := ReceiverOptions{
+		HeartbeatTimeout: 5 * time.Second,
+	}
+	r := NewReceiverWithOpts(log.WithContext(context.Background()), opts)
 	defer r.Stop()
 
 	select {
@@ -38,18 +25,10 @@ func TestNewReceiver(t *testing.T) {
 }
 
 func TestNewReceiverNotReceive(t *testing.T) {
-	testMutex.Lock()
-	originalTimeout := heartbeatTimeout
-	heartbeatTimeout = 1 * time.Second
-	testMutex.Unlock()
-
-	defer func() {
-		testMutex.Lock()
-		heartbeatTimeout = originalTimeout
-		testMutex.Unlock()
-	}()
-
-	r := NewReceiver(log.WithContext(context.Background()))
+	opts := ReceiverOptions{
+		HeartbeatTimeout: 1 * time.Second,
+	}
+	r := NewReceiverWithOpts(log.WithContext(context.Background()), opts)
 	defer r.Stop()
 
 	select {
@@ -61,18 +40,10 @@ func TestNewReceiverNotReceive(t *testing.T) {
 }
 
 func TestNewReceiverAck(t *testing.T) {
-	testMutex.Lock()
-	originalTimeout := heartbeatTimeout
-	heartbeatTimeout = 2 * time.Second
-	testMutex.Unlock()
-
-	defer func() {
-		testMutex.Lock()
-		heartbeatTimeout = originalTimeout
-		testMutex.Unlock()
-	}()
-
-	r := NewReceiver(log.WithContext(context.Background()))
+	opts := ReceiverOptions{
+		HeartbeatTimeout: 2 * time.Second,
+	}
+	r := NewReceiverWithOpts(log.WithContext(context.Background()), opts)
 	defer r.Stop()
 
 	r.Heartbeat()
@@ -97,30 +68,19 @@ func TestReceiverHealthCheckAttemptThreshold(t *testing.T) {
 
 	for _, tc := range testsCases {
 		t.Run(tc.name, func(t *testing.T) {
-			testMutex.Lock()
-			originalInterval := healthCheckInterval
-			originalTimeout := heartbeatTimeout
-			healthCheckInterval = 1 * time.Second
-			heartbeatTimeout = healthCheckInterval + 500*time.Millisecond
-			testMutex.Unlock()
-
-			defer func() {
-				testMutex.Lock()
-				healthCheckInterval = originalInterval
-				heartbeatTimeout = originalTimeout
-				testMutex.Unlock()
-			}()
-			//nolint:tenv
-			os.Setenv(defaultAttemptThresholdEnv, fmt.Sprintf("%d", tc.threshold))
-			defer os.Unsetenv(defaultAttemptThresholdEnv)
-
-			receiver := NewReceiver(log.WithField("test_name", tc.name))
-
-			testTimeout := heartbeatTimeout*time.Duration(tc.threshold) + healthCheckInterval
+			healthCheckInterval := 1 * time.Second
+
+			opts := ReceiverOptions{
+				HeartbeatTimeout: healthCheckInterval + 500*time.Millisecond,
+				AttemptThreshold: tc.threshold,
+			}
+
+			receiver := NewReceiverWithOpts(log.WithField("test_name", tc.name), opts)
+
+			testTimeout := opts.HeartbeatTimeout*time.Duration(tc.threshold) + healthCheckInterval
 
 			if tc.resetCounterOnce {
 				receiver.Heartbeat()
-				t.Logf("reset counter once")
 			}
 
 			select {
@@ -134,7 +94,6 @@ func TestReceiverHealthCheckAttemptThreshold(t *testing.T) {
 				}
 				t.Fatalf("should have timed out before %s", testTimeout)
 			}
-
 		})
 	}
 }