nsq_scaler.go

package scalers

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"sync"

	"github.com/go-logr/logr"
	"github.com/kedacore/keda/v2/pkg/scalers/scalersconfig"
	kedautil "github.com/kedacore/keda/v2/pkg/util"
	v2 "k8s.io/api/autoscaling/v2"
	"k8s.io/metrics/pkg/apis/external_metrics"
)

type nsqScaler struct {
	metricType v2.MetricTargetType
	metadata   nsqMetadata
	httpClient *http.Client
	logger     logr.Logger
}

type nsqMetadata struct {
	NSQLookupdHTTPAddresses  []string `keda:"name=nsqLookupdHTTPAddresses,  order=triggerMetadata;resolvedEnv"`
	Topic                    string   `keda:"name=topic,                    order=triggerMetadata;resolvedEnv"`
	Channel                  string   `keda:"name=channel,                  order=triggerMetadata;resolvedEnv"`
	DepthThreshold           int64    `keda:"name=depthThreshold,           order=triggerMetadata;resolvedEnv, default=10"`
	ActivationDepthThreshold int64    `keda:"name=activationDepthThreshold, order=triggerMetadata;resolvedEnv, default=0"`

	triggerIndex int
}

const (
	nsqMetricType = "External"
)

func NewNSQScaler(config *scalersconfig.ScalerConfig) (Scaler, error) {
	metricType, err := GetMetricTargetType(config)
	if err != nil {
		return nil, fmt.Errorf("error getting scaler metric type: %w", err)
	}

	logger := InitializeLogger(config, "nsq_scaler")

	nsqMetadata, err := parseNSQMetadata(config)
	if err != nil {
		return nil, fmt.Errorf("error parsing NSQ metadata: %w", err)
	}

	return &nsqScaler{
		metricType: metricType,
		metadata:   nsqMetadata,
		httpClient: kedautil.CreateHTTPClient(config.GlobalHTTPTimeout, true),
		logger:     logger,
	}, nil
}

func (m nsqMetadata) Validate() error {
	if len(m.NSQLookupdHTTPAddresses) == 0 {
		return fmt.Errorf("no nsqLookupdHTTPAddresses given")
	}

	if m.Topic == "" {
		return fmt.Errorf("no topic given")
	}

	if m.Channel == "" {
		return fmt.Errorf("no channel given")
	}

	if m.DepthThreshold <= 0 {
		return fmt.Errorf("depthThreshold must be a positive integer")
	}

	if m.ActivationDepthThreshold < 0 {
		return fmt.Errorf("activationDepthThreshold must be greater than or equal to 0")
	}

	return nil
}

func parseNSQMetadata(config *scalersconfig.ScalerConfig) (nsqMetadata, error) {
	meta := nsqMetadata{triggerIndex: config.TriggerIndex}
	if err := config.TypedConfig(&meta); err != nil {
		return meta, fmt.Errorf("error parsing nsq metadata: %w", err)
	}

	return meta, nil
}

func (s nsqScaler) GetMetricsAndActivity(ctx context.Context, metricName string) ([]external_metrics.ExternalMetricValue, bool, error) {
	depth, err := s.getTopicChannelDepth()

	if err != nil {
		return []external_metrics.ExternalMetricValue{}, false, err
	}

	s.logger.Info("GetMetricsAndActivity", "metricName", metricName, "depth", depth)

	metric := GenerateMetricInMili(metricName, float64(depth))

	return []external_metrics.ExternalMetricValue{metric}, depth > s.metadata.ActivationDepthThreshold, nil
}

func (s nsqScaler) getTopicChannelDepth() (int64, error) {
	nsqdHosts, err := s.getTopicProducers(s.metadata.Topic)
	if err != nil {
		return -1, fmt.Errorf("error getting nsqd hosts: %w", err)
	}

	if len(nsqdHosts) == 0 {
		s.logger.Info("no nsqd hosts found for topic", "topic", s.metadata.Topic)
		return 0, nil
	}

	depth, err := s.aggregateDepth(nsqdHosts, s.metadata.Topic, s.metadata.Channel)
	if err != nil {
		return -1, fmt.Errorf("error getting topic/channel depth: %w", err)
	}

	return depth, nil
}

func (s nsqScaler) GetMetricSpecForScaling(ctx context.Context) []v2.MetricSpec {
	metricName := fmt.Sprintf("nsq-%s-%s", s.metadata.Topic, s.metadata.Channel)

	externalMetric := &v2.ExternalMetricSource{
		Metric: v2.MetricIdentifier{
			Name: GenerateMetricNameWithIndex(s.metadata.triggerIndex, kedautil.NormalizeString(metricName)),
		},
		Target: GetMetricTarget(s.metricType, s.metadata.DepthThreshold),
	}
	metricSpec := v2.MetricSpec{External: externalMetric, Type: nsqMetricType}
	return []v2.MetricSpec{metricSpec}
}

func (s nsqScaler) Close(ctx context.Context) error {
	if s.httpClient != nil {
		s.httpClient.CloseIdleConnections()
	}
	return nil
}

type lookupResponse struct {
	Producers []struct {
		HttpPort         int    `json:"http_port"`
		BroadcastAddress string `json:"broadcast_address"`
	}
}

type lookupResult struct {
	host           string
	lookupResponse *lookupResponse
	err            error
}

func (s *nsqScaler) getTopicProducers(topic string) ([]string, error) {
	var wg sync.WaitGroup
	resultCh := make(chan lookupResult, len(s.metadata.NSQLookupdHTTPAddresses))

	for _, host := range s.metadata.NSQLookupdHTTPAddresses {
		wg.Add(1)
		go func(host string, topic string) {
			defer wg.Done()
			resp, err := s.getLookup(host, topic)
			resultCh <- lookupResult{host, resp, err}
		}(host, topic)
	}

	wg.Wait()
	close(resultCh)

	var nsqdHostMap = make(map[string]bool)
	for result := range resultCh {
		if result.err != nil {
			return nil, fmt.Errorf("error getting lookup from host '%s': %w", result.host, result.err)
		}

		if result.lookupResponse == nil {
			// topic is not found on a single nsqlookupd host, it may exist on another
			continue
		}

		for _, producer := range result.lookupResponse.Producers {
			nsqdHost := fmt.Sprintf("%s:%d", producer.BroadcastAddress, producer.HttpPort)
			nsqdHostMap[nsqdHost] = true
		}
	}

	var nsqdHosts []string
	for nsqdHost := range nsqdHostMap {
		nsqdHosts = append(nsqdHosts, nsqdHost)
	}

	return nsqdHosts, nil
}

func (s *nsqScaler) getLookup(host string, topic string) (*lookupResponse, error) {
	req, err := http.NewRequest("GET", fmt.Sprintf("http://%s/%s", host, "lookup"), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/json; charset=utf-8")

	params := url.Values{"topic": {topic}}
	req.URL.RawQuery = params.Encode()

	resp, err := s.httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusNotFound {
		return nil, nil
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("unexpected status code '%s'", resp.Status)
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var lookupResponse lookupResponse
	err = json.Unmarshal(body, &lookupResponse)
	if err != nil {
		return nil, err
	}

	return &lookupResponse, nil
}

type statsResponse struct {
	Topics []struct {
		TopicName string `json:"topic_name"`
		Depth     int64  `json:"depth"`
		Channels  []struct {
			ChannelName string `json:"channel_name"`
			Depth       int64  `json:"depth"`  // num messages in the queue (mem + disk)
			Paused      bool   `json:"paused"` // if paused, consumers will not receive messages
		}
	}
}

type statsResult struct {
	host          string
	statsResponse *statsResponse
	err           error
}

func (s *nsqScaler) aggregateDepth(nsqdHosts []string, topic string, channel string) (int64, error) {
	wg := sync.WaitGroup{}
	resultCh := make(chan statsResult, len(nsqdHosts))

	for _, host := range nsqdHosts {
		wg.Add(1)
		go func(host string, topic string) {
			defer wg.Done()
			resp, err := s.getStats(host, topic)
			resultCh <- statsResult{host, resp, err}
		}(host, topic)
	}

	wg.Wait()
	close(resultCh)

	var depth int64
	for result := range resultCh {
		if result.err != nil {
			return -1, fmt.Errorf("error getting stats from host '%s': %w", result.host, result.err)
		}

		for _, t := range result.statsResponse.Topics {
			if t.TopicName != topic {
				// this should never happen as we make the /stats call with the "topic" param
				continue
			}

			if len(t.Channels) == 0 {
				// topic exists with no channels, but there are messages in the topic -> we should still scale to bootstrap
				s.logger.Info("no channels exist for topic", "topic", topic, "channel", channel, "host", result.host)
				depth += t.Depth
				continue
			}

			channelExists := false
			for _, ch := range t.Channels {
				if ch.ChannelName != channel {
					continue
				}
				channelExists = true
				if ch.Paused {
					// if it's paused on a single nsqd host, it's depth should not go into the aggregate
					// meaning if paused on all nsqd hosts => depth == 0
					s.logger.Info("channel is paused", "topic", topic, "channel", channel, "host", result.host)
					continue
				}
				depth += ch.Depth
			}
			if !channelExists {
				// topic exists with channels, but not the one in question - fallback to topic depth
				s.logger.Info("channel does not exist for topic", "topic", topic, "channel", channel, "host", result.host)
				depth += t.Depth
			}
		}
	}

	return depth, nil
}

func (s *nsqScaler) getStats(host string, topic string) (*statsResponse, error) {
	req, err := http.NewRequest("GET", fmt.Sprintf("http://%s/%s", host, "stats"), nil)
	if err != nil {
		return nil, err
	}

	// "channel" is a query param as well, but if used and the channel does not exist
	// we do not receive any stats for the existing topic
	params := url.Values{
		"format":          {"json"},
		"include_clients": {"false"},
		"include_mem":     {"false"},
		"topic":           {topic},
	}
	req.URL.RawQuery = params.Encode()

	resp, err := s.httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("unexpected status code '%s'", resp.Status)
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var statsResponse statsResponse
	err = json.Unmarshal(body, &statsResponse)
	if err != nil {
		return nil, err
	}

	return &statsResponse, nil
}