[relay] Add health check attempt threshold (#2609)

* Add health check attempt threshold for receiver

* Add health check attempt threshold for sender

relay/healthcheck/receiver_test.go

@@ -1,13 +1,18 @@
 package healthcheck
 
 import (
+	"context"
+	"fmt"
+	"os"
 	"testing"
 	"time"
+
+	log "github.com/sirupsen/logrus"
 )
 
 func TestNewReceiver(t *testing.T) {
 	heartbeatTimeout = 5 * time.Second
-	r := NewReceiver()
+	r := NewReceiver(log.WithContext(context.Background()))
 
 	select {
 	case <-r.OnTimeout:
@@ -19,7 +24,7 @@ func TestNewReceiver(t *testing.T) {
 
 func TestNewReceiverNotReceive(t *testing.T) {
 	heartbeatTimeout = 1 * time.Second
-	r := NewReceiver()
+	r := NewReceiver(log.WithContext(context.Background()))
 
 	select {
 	case <-r.OnTimeout:
@@ -30,7 +35,7 @@ func TestNewReceiverNotReceive(t *testing.T) {
 
 func TestNewReceiverAck(t *testing.T) {
 	heartbeatTimeout = 2 * time.Second
-	r := NewReceiver()
+	r := NewReceiver(log.WithContext(context.Background()))
 
 	r.Heartbeat()
 
@@ -40,3 +45,53 @@ func TestNewReceiverAck(t *testing.T) {
 	case <-time.After(3 * time.Second):
 	}
 }
+
+func TestReceiverHealthCheckAttemptThreshold(t *testing.T) {
+	testsCases := []struct {
+		name             string
+		threshold        int
+		resetCounterOnce bool
+	}{
+		{"Default attempt threshold", defaultAttemptThreshold, false},
+		{"Custom attempt threshold", 3, false},
+		{"Should reset threshold once", 2, true},
+	}
+
+	for _, tc := range testsCases {
+		t.Run(tc.name, func(t *testing.T) {
+			originalInterval := healthCheckInterval
+			originalTimeout := heartbeatTimeout
+			healthCheckInterval = 1 * time.Second
+			heartbeatTimeout = healthCheckInterval + 500*time.Millisecond
+			defer func() {
+				healthCheckInterval = originalInterval
+				heartbeatTimeout = originalTimeout
+			}()
+			//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
+
+			if tc.resetCounterOnce {
+				receiver.Heartbeat()
+				t.Logf("reset counter once")
+			}
+
+			select {
+			case <-receiver.OnTimeout:
+				if tc.resetCounterOnce {
+					t.Fatalf("should not have timed out before %s", testTimeout)
+				}
+			case <-time.After(testTimeout):
+				if tc.resetCounterOnce {
+					return
+				}
+				t.Fatalf("should have timed out before %s", testTimeout)
+			}
+
+		})
+	}
+}