tests

Signed-off-by: alanprot <[email protected]>

Author: alanprot

pkg/cortex/modules.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	"github.com/cortexproject/cortex/pkg/util/grpcclient"
 	"net/http"
 
 	"github.com/go-kit/log"
@@ -49,6 +48,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/scheduler"
 	"github.com/cortexproject/cortex/pkg/storage/bucket"
 	"github.com/cortexproject/cortex/pkg/storegateway"
+	"github.com/cortexproject/cortex/pkg/util/grpcclient"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 	"github.com/cortexproject/cortex/pkg/util/modules"
 	"github.com/cortexproject/cortex/pkg/util/runtimeconfig"

pkg/util/grpcclient/health_check.go

@@ -53,12 +53,15 @@ type HealthCheckInterceptors struct {
 
 	sync.RWMutex
 	activeInstances map[string]*healthCheckEntry
+
+	healthClientFactory func(cc grpc.ClientConnInterface) grpc_health_v1.HealthClient
 }
 
 func NewHealthCheckInterceptors(logger log.Logger) *HealthCheckInterceptors {
 	h := &HealthCheckInterceptors{
-		logger:          logger,
-		activeInstances: make(map[string]*healthCheckEntry),
+		logger:              logger,
+		healthClientFactory: grpc_health_v1.NewHealthClient,
+		activeInstances:     make(map[string]*healthCheckEntry),
 	}
 
 	h.Service = services.
@@ -105,7 +108,7 @@ func (h *HealthCheckInterceptors) iteration(ctx context.Context) error {
 			return err
 		}
 		conn, err := grpc.NewClient(instance.address, dialOpts...)
-		c := grpc_health_v1.NewHealthClient(conn)
+		c := h.healthClientFactory(conn)
 		if err != nil {
 			return err
 		}
@@ -120,6 +123,7 @@ func (h *HealthCheckInterceptors) iteration(ctx context.Context) error {
 		if time.Since(instance.lastCheckTime.Load()) < instance.clientConfig.HealthCheckConfig.Interval {
 			continue
 		}
+		instance.lastCheckTime.Store(time.Now())
 
 		go func(i *healthCheckEntry) {
 			i.recordHealth(healthCheck(c, i.clientConfig.HealthCheckConfig.Timeout))

pkg/util/grpcclient/health_check_test.go

@@ -0,0 +1,85 @@
+package grpcclient
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	cortex_testutil "github.com/cortexproject/cortex/pkg/util/test"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/health/grpc_health_v1"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	utillog "github.com/cortexproject/cortex/pkg/util/log"
+)
+
+type healthClientMock struct {
+	grpc_health_v1.HealthClient
+	err error
+}
+
+func (h *healthClientMock) Check(ctx context.Context, in *grpc_health_v1.HealthCheckRequest, opts ...grpc.CallOption) (*grpc_health_v1.HealthCheckResponse, error) {
+	return &grpc_health_v1.HealthCheckResponse{
+		Status: grpc_health_v1.HealthCheckResponse_SERVING,
+	}, h.err
+}
+
+func TestNewHealthCheckInterceptors(t *testing.T) {
+	i := NewHealthCheckInterceptors(utillog.Logger)
+	hMock := &healthClientMock{
+		err: fmt.Errorf("some error"),
+	}
+	cfg := Config{
+		HealthCheckConfig: HealthCheckConfig{
+			UnhealthyThreshold: 2,
+			Interval:           0,
+			Timeout:            time.Second,
+		},
+	}
+	i.healthClientFactory = func(cc grpc.ClientConnInterface) grpc_health_v1.HealthClient {
+		return hMock
+	}
+
+	ui := i.UnaryHealthCheckInterceptor(cfg)
+	ccUnhealthy, err := grpc.NewClient("localhost:999", grpc.WithTransportCredentials(insecure.NewCredentials()))
+	ccHealthy, err := grpc.NewClient("localhost:111", grpc.WithTransportCredentials(insecure.NewCredentials()))
+	require.NoError(t, err)
+	invokedMap := map[string]int{}
+
+	invoker := func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
+		invokedMap[cc.Target()]++
+		return nil
+	}
+
+	//Should allow first call
+	require.NoError(t, ui(context.Background(), "", struct{}{}, struct{}{}, ccUnhealthy, invoker))
+
+	// first health check
+	i.iteration(context.Background())
+
+	//Should second first call
+	require.NoError(t, ui(context.Background(), "", struct{}{}, struct{}{}, ccUnhealthy, invoker))
+
+	require.Equal(t, invokedMap["localhost:999"], 2)
+
+	// Second Healthcheck -> should mark as unhealthy
+	i.iteration(context.Background())
+
+	cortex_testutil.Poll(t, time.Second, true, func() interface{} {
+		return errors.Is(ui(context.Background(), "", struct{}{}, struct{}{}, ccUnhealthy, invoker), unhealthyErr)
+	})
+
+	// Other instances should remain healthy
+	require.NoError(t, ui(context.Background(), "", struct{}{}, struct{}{}, ccHealthy, invoker))
+
+	// Should mark the instance back to healthy
+	hMock.err = nil
+	i.iteration(context.Background())
+	cortex_testutil.Poll(t, time.Second, true, func() interface{} {
+		return ui(context.Background(), "", struct{}{}, struct{}{}, ccUnhealthy, invoker) == nil
+	})
+
+}