Querying the stream consumer's leader node when using a cluster for the accurate count

rayjanoka · rayjanoka · commit 3c7505431b81 · 2022-11-13T00:21:03.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -36,6 +36,10 @@ To learn more about active deprecations, we recommend checking [GitHub Discussio
 - [v1.1.0](#v110)
 - [v1.0.0](#v100)
 
+## Unreleased
+
+- **JetStream:** JetStream scaler query the stream consumer's leader pod when clustered ([#2391](https://github.com/kedacore/keda/issues/2391))
+
 ### New
 
 - **General**: Expand Prometheus metric with label "ScalerName" to distinguish different triggers. The scaleName is defined per Trigger.Name ([#3588](https://github.com/kedacore/keda/issues/3588))
diff --git a/pkg/scalers/nats_jetstream_scaler.go b/pkg/scalers/nats_jetstream_scaler.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"net/url"
 	"strconv"
 
 	"github.com/go-logr/logr"
@@ -26,32 +27,40 @@ const (
 )
 
 type natsJetStreamScaler struct {
-	metricType v2.MetricTargetType
-	stream     *streamDetail
-	metadata   natsJetStreamMetadata
-	httpClient *http.Client
-	logger     logr.Logger
+	metricType   v2.MetricTargetType
+	stream       *streamDetail
+	metadata     natsJetStreamMetadata
+	httpClient   *http.Client
+	logger       logr.Logger
+	leaderStream *streamDetail
 }
 
 type natsJetStreamMetadata struct {
-	monitoringEndpoint     string
 	account                string
 	stream                 string
 	consumer               string
+	leaderName             string
+	monitoringURL          string
 	lagThreshold           int64
 	activationLagThreshold int64
+	clusterSize            int
 	scalerIndex            int
 }
 
 type jetStreamEndpointResponse struct {
-	Accounts []accountDetail `json:"account_details"`
+	Accounts    []accountDetail `json:"account_details"`
+	MetaCluster metaCluster     `json:"meta_cluster"`
 }
 
 type accountDetail struct {
 	Name    string          `json:"name"`
 	Streams []*streamDetail `json:"stream_detail"`
 }
 
+type metaCluster struct {
+	ClusterSize int `json:"cluster_size"`
+}
+
 type streamDetail struct {
 	Name      string           `json:"name"`
 	Config    streamConfig     `json:"config"`
@@ -77,6 +86,11 @@ type consumerDetail struct {
 	NumPending     int                    `json:"num_pending"`
 	Config         consumerConfig         `json:"config"`
 	DeliveryStatus consumerDeliveryStatus `json:"delivery"`
+	Cluster        consumerCluster        `json:"cluster"`
+}
+
+type consumerCluster struct {
+	Leader string `json:"leader"`
 }
 
 type consumerConfig struct {
@@ -158,61 +172,106 @@ func parseNATSJetStreamMetadata(config *ScalerConfig) (natsJetStreamMetadata, er
 			return meta, fmt.Errorf("useHTTPS parsing error %s", err.Error())
 		}
 	}
-	meta.monitoringEndpoint = getNATSJetStreamEndpoint(useHTTPS, natsServerEndpoint, meta.account)
+	meta.monitoringURL = getNATSJetStreamEndpoint(useHTTPS, natsServerEndpoint, meta.account)
 
 	return meta, nil
 }
 
 func getNATSJetStreamEndpoint(useHTTPS bool, natsServerEndpoint string, account string) string {
-	protocol := natsHTTPProtocol
+	scheme := natsHTTPProtocol
 	if useHTTPS {
-		protocol = natsHTTPSProtocol
+		scheme = natsHTTPSProtocol
 	}
+	return fmt.Sprintf("%s://%s/jsz?acc=%s&consumers=true&config=true", scheme, natsServerEndpoint, account)
+}
 
-	return fmt.Sprintf("%s://%s/jsz?acc=%s&consumers=true&config=true", protocol, natsServerEndpoint, account)
+func (s *natsJetStreamScaler) getNATSJetStreamLeaderEndpoint() (string, error) {
+	jsURL, err := url.Parse(s.metadata.monitoringURL)
+	if err != nil {
+		s.logger.Error(err, "unable to parse monitoring URL to create leader URL", "natsServerMonitoringURL", s.metadata.monitoringURL)
+		return "", err
+	}
+
+	return fmt.Sprintf("%s://%s.%s%s?%s", jsURL.Scheme, s.metadata.leaderName, jsURL.Host, jsURL.Path, jsURL.RawQuery), nil
 }
 
-func (s *natsJetStreamScaler) IsActive(ctx context.Context) (bool, error) {
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.metadata.monitoringEndpoint, nil)
+func (s *natsJetStreamScaler) getNATSJetstreamMonitoringData(ctx context.Context, natsJetStreamEndpoint string) error {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, natsJetStreamEndpoint, nil)
 	if err != nil {
-		return false, err
+		return err
 	}
 
 	resp, err := s.httpClient.Do(req)
 	if err != nil {
-		s.logger.Error(err, "unable to access NATS JetStream monitoring endpoint", "natsServerMonitoringEndpoint", s.metadata.monitoringEndpoint)
-		return false, err
+		s.logger.Error(err, "unable to access NATS JetStream monitoring endpoint", "natsServerMonitoringEndpoint", natsJetStreamEndpoint)
+		return err
 	}
 
 	defer resp.Body.Close()
 	var jsAccountResp jetStreamEndpointResponse
 	if err = json.NewDecoder(resp.Body).Decode(&jsAccountResp); err != nil {
-		s.logger.Error(err, "unable to decode JetStream account response")
-		return false, err
+		s.logger.Error(err, "unable to decode NATS JetStream account details")
+		return err
 	}
 
+	s.metadata.clusterSize = jsAccountResp.MetaCluster.ClusterSize
+
 	// Find and assign the stream that we are looking for.
-	for _, account := range jsAccountResp.Accounts {
-		if account.Name == s.metadata.account {
-			for _, stream := range account.Streams {
+	for _, jsAccount := range jsAccountResp.Accounts {
+		if jsAccount.Name == s.metadata.account {
+			for _, stream := range jsAccount.Streams {
 				if stream.Name == s.metadata.stream {
 					s.stream = stream
 				}
+
+				for _, consumer := range stream.Consumers {
+					if consumer.Name == s.metadata.consumer {
+						s.metadata.leaderName = consumer.Cluster.Leader
+					}
+				}
 			}
 		}
 	}
+
+	return nil
+}
+
+func (s *natsJetStreamScaler) IsActive(ctx context.Context) (bool, error) {
+	err := s.getNATSJetstreamMonitoringData(ctx, s.metadata.monitoringURL)
+	if err != nil {
+		return false, err
+	}
+
+	// Query the consumer leader pod, it has the accurate count.
+	if s.metadata.clusterSize > 1 {
+		monitoringLeaderEndpoint, err := s.getNATSJetStreamLeaderEndpoint()
+		if err != nil {
+			return false, err
+		}
+
+		err = s.getNATSJetstreamMonitoringData(ctx, monitoringLeaderEndpoint)
+		if err != nil {
+			return false, err
+		}
+	}
+
 	return s.getMaxMsgLag() > s.metadata.activationLagThreshold, nil
 }
 
 func (s *natsJetStreamScaler) getMaxMsgLag() int64 {
 	consumerName := s.metadata.consumer
 
-	for _, consumer := range s.stream.Consumers {
+	stream := s.stream
+	if s.leaderStream != nil {
+		stream = s.leaderStream
+	}
+
+	for _, consumer := range stream.Consumers {
 		if consumer.Name == consumerName {
 			return int64(consumer.NumPending + consumer.NumAckPending)
 		}
 	}
-	return s.stream.State.LastSequence
+	return stream.State.LastSequence
 }
 
 func (s *natsJetStreamScaler) GetMetricSpecForScaling(context.Context) []v2.MetricSpec {
@@ -230,32 +289,21 @@ func (s *natsJetStreamScaler) GetMetricSpecForScaling(context.Context) []v2.Metr
 }
 
 func (s *natsJetStreamScaler) GetMetrics(ctx context.Context, metricName string, metricSelector labels.Selector) ([]external_metrics.ExternalMetricValue, error) {
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.metadata.monitoringEndpoint, nil)
-	if err != nil {
-		return nil, err
-	}
-
-	resp, err := s.httpClient.Do(req)
+	err := s.getNATSJetstreamMonitoringData(ctx, s.metadata.monitoringURL)
 	if err != nil {
-		s.logger.Error(err, "unable to access NATS JetStream monitoring endpoint", "natsServerMonitoringEndpoint", s.metadata.monitoringEndpoint)
 		return []external_metrics.ExternalMetricValue{}, err
 	}
 
-	defer resp.Body.Close()
-	var jsAccountResp jetStreamEndpointResponse
-	if err = json.NewDecoder(resp.Body).Decode(&jsAccountResp); err != nil {
-		s.logger.Error(err, "unable to decode JetStream account details")
-		return []external_metrics.ExternalMetricValue{}, err
-	}
+	// Query the consumer leader pod, it has the accurate count.
+	if s.metadata.clusterSize > 1 {
+		monitoringLeaderEndpoint, err := s.getNATSJetStreamLeaderEndpoint()
+		if err != nil {
+			return []external_metrics.ExternalMetricValue{}, err
+		}
 
-	// Find and assign the stream that we are looking for.
-	for _, account := range jsAccountResp.Accounts {
-		if account.Name == s.metadata.account {
-			for _, stream := range account.Streams {
-				if stream.Name == s.metadata.stream {
-					s.stream = stream
-				}
-			}
+		err = s.getNATSJetstreamMonitoringData(ctx, monitoringLeaderEndpoint)
+		if err != nil {
+			return []external_metrics.ExternalMetricValue{}, err
 		}
 	}
 
diff --git a/pkg/scalers/nats_jetstream_scaler_test.go b/pkg/scalers/nats_jetstream_scaler_test.go
@@ -67,13 +67,13 @@ func TestNATSJetStreamGetMetricSpecForScaling(t *testing.T) {
 		if err != nil {
 			t.Fatal("Could not parse metadata:", err)
 		}
-		mockStanScaler := natsJetStreamScaler{
+		mockJetStreamScaler := natsJetStreamScaler{
 			stream:     nil,
 			metadata:   meta,
 			httpClient: http.DefaultClient,
 		}
 
-		metricSpec := mockStanScaler.GetMetricSpecForScaling(ctx)
+		metricSpec := mockJetStreamScaler.GetMetricSpecForScaling(ctx)
 		metricName := metricSpec[0].External.Metric.Name
 		if metricName != testData.name {
 			t.Error("Wrong External metric source name:", metricName)
