azure_pipelines_scaler.go

package scalers

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

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
	"github.com/go-logr/logr"
	v2 "k8s.io/api/autoscaling/v2"
	"k8s.io/metrics/pkg/apis/external_metrics"

	kedav1alpha1 "github.com/kedacore/keda/v2/apis/keda/v1alpha1"
	"github.com/kedacore/keda/v2/pkg/scalers/azure"
	"github.com/kedacore/keda/v2/pkg/scalers/scalersconfig"
	kedautil "github.com/kedacore/keda/v2/pkg/util"
)

type JobRequests struct {
	Count int          `json:"count"`
	Value []JobRequest `json:"value"`
}

const (
	// "499b84ac-1321-427f-aa17-267ca6975798" is the azure id for DevOps resource
	// https://learn.microsoft.com/en-gb/azure/devops/integrate/get-started/authentication/service-principal-managed-identity?view=azure-devops
	devopsResource = "499b84ac-1321-427f-aa17-267ca6975798/.default"
)

type JobRequest struct {
	RequestID     int       `json:"requestId"`
	QueueTime     time.Time `json:"queueTime"`
	AssignTime    time.Time `json:"assignTime"`
	ReceiveTime   time.Time `json:"receiveTime"`
	LockedUntil   time.Time `json:"lockedUntil"`
	ServiceOwner  string    `json:"serviceOwner"`
	HostID        string    `json:"hostId"`
	Result        *string   `json:"result"`
	ScopeID       string    `json:"scopeId"`
	PlanType      string    `json:"planType"`
	PlanID        string    `json:"planId"`
	JobID         string    `json:"jobId"`
	Demands       []string  `json:"demands"`
	ReservedAgent *struct {
		Links struct {
			Self struct {
				Href string `json:"href"`
			} `json:"self"`
			Web struct {
				Href string `json:"href"`
			} `json:"web"`
		} `json:"_links"`
		ID                int    `json:"id"`
		Name              string `json:"name"`
		Version           string `json:"version"`
		OsDescription     string `json:"osDescription"`
		Enabled           bool   `json:"enabled"`
		Status            string `json:"status"`
		ProvisioningState string `json:"provisioningState"`
		AccessPoint       string `json:"accessPoint"`
	} `json:"reservedAgent,omitempty"`
	Definition struct {
		Links struct {
			Web struct {
				Href string `json:"href"`
			} `json:"web"`
			Self struct {
				Href string `json:"href"`
			} `json:"self"`
		} `json:"_links"`
		ID   int    `json:"id"`
		Name string `json:"name"`
	} `json:"definition"`
	Owner struct {
		Links struct {
			Web struct {
				Href string `json:"href"`
			} `json:"web"`
			Self struct {
				Href string `json:"href"`
			} `json:"self"`
		} `json:"_links"`
		ID   int    `json:"id"`
		Name string `json:"name"`
	} `json:"owner"`
	Data struct {
		ParallelismTag string `json:"ParallelismTag"`
		IsScheduledKey string `json:"IsScheduledKey"`
	} `json:"data"`
	PoolID          int    `json:"poolId"`
	OrchestrationID string `json:"orchestrationId"`
	Priority        int    `json:"priority"`
	MatchedAgents   *[]struct {
		Links struct {
			Self struct {
				Href string `json:"href"`
			} `json:"self"`
			Web struct {
				Href string `json:"href"`
			} `json:"web"`
		} `json:"_links"`
		ID                int    `json:"id"`
		Name              string `json:"name"`
		Version           string `json:"version"`
		Enabled           bool   `json:"enabled"`
		Status            string `json:"status"`
		ProvisioningState string `json:"provisioningState"`
	} `json:"matchedAgents,omitempty"`
}

type azurePipelinesPoolNameResponse struct {
	Value []struct {
		ID int `json:"id"`
	} `json:"value"`
}

type azurePipelinesPoolIDResponse struct {
	ID int `json:"id"`
}

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

type azurePipelinesMetadata struct {
	OrganizationURL                      string `keda:"name=organizationURL, order=resolvedEnv;authParams;triggerMetadata"`
	OrganizationName                     string
	PersonalAccessToken                  string `keda:"name=personalAccessToken, order=authParams;resolvedEnv, optional"`
	authContext                          authContext
	Parent                               string `keda:"name=parent, order=triggerMetadata, optional"`
	Demands                              string `keda:"name=demands, order=triggerMetadata, optional"`
	PoolName                             string `keda:"name=poolName, order=triggerMetadata, optional"`
	PoolID                               int    `keda:"name=poolID, order=triggerMetadata, optional"`
	TargetPipelinesQueueLength           int64  `keda:"name=targetPipelinesQueueLength, order=triggerMetadata, default=1"`
	ActivationTargetPipelinesQueueLength int64  `keda:"name=activationTargetPipelinesQueueLength, order=triggerMetadata, default=0"`
	JobsToFetch                          int64  `keda:"name=jobsToFetch, order=triggerMetadata, default=250"`
	FetchUnfinishedJobsOnly              bool   `keda:"name=fetchUnfinishedJobsOnly, order=triggerMetadata, default=false"`
	ScaleOnInFlight                      bool   `keda:"name=scaleOnInFlight, order=triggerMetadata, default=true"`
	triggerIndex                         int
	RequireAllDemands                    bool `keda:"name=requireAllDemands, order=triggerMetadata, default=false"`
	RequireAllDemandsAndIgnoreOthers     bool `keda:"name=requireAllDemandsAndIgnoreOthers, order=triggerMetadata, default=false"`
	CaseInsensitiveDemandsProcessing     bool `keda:"name=caseInsensitiveDemandsProcessing, order=triggerMetadata, default=false"`
}

type authContext struct {
	credential azcore.TokenCredential
	pat        string
	token      *azcore.AccessToken
	authType   azurePipelinesAuthType
}

type azurePipelinesAuthType string

const (
	azurePipelinesAuthTypePAT              azurePipelinesAuthType = "pat"
	azurePipelinesAuthTypeServicePrincipal azurePipelinesAuthType = "service-principal"
	azurePipelinesAuthTypeWorkloadIdentity azurePipelinesAuthType = "workload-identity"
)

func (a *azurePipelinesMetadata) Validate() error {
	if val := a.OrganizationURL[strings.LastIndex(a.OrganizationURL, "/")+1:]; val != "" {
		a.OrganizationName = a.OrganizationURL[strings.LastIndex(a.OrganizationURL, "/")+1:]
	} else {
		return fmt.Errorf("failed to extract organization name from organizationURL")
	}
	return nil
}

// NewAzurePipelinesScaler creates a new AzurePipelinesScaler
func NewAzurePipelinesScaler(ctx context.Context, config *scalersconfig.ScalerConfig) (Scaler, error) {
	httpClient := kedautil.CreateHTTPClient(config.GlobalHTTPTimeout, false)

	logger := InitializeLogger(config, "azure_pipelines_scaler")
	metricType, err := GetMetricTargetType(config)
	if err != nil {
		return nil, fmt.Errorf("error getting scaler metric type: %w", err)
	}

	meta, err := parseAzurePipelinesMetadata(ctx, logger, config, httpClient)
	if err != nil {
		return nil, fmt.Errorf("error parsing azure Pipelines metadata: %w", err)
	}

	return &azurePipelinesScaler{
		metricType: metricType,
		metadata:   meta,
		httpClient: httpClient,
		logger:     logger,
	}, nil
}

func configureAzurePipelinesAuth(logger logr.Logger, config *scalersconfig.ScalerConfig, meta *azurePipelinesMetadata) error {
	if meta.PersonalAccessToken != "" {
		meta.authContext.pat = strings.TrimSuffix(meta.PersonalAccessToken, "\n")
		meta.authContext.authType = azurePipelinesAuthTypePAT
		return nil
	}

	switch config.PodIdentity.Provider {
	case "", kedav1alpha1.PodIdentityProviderNone:
		if azure.IsServicePrincipalAuthConfigured(config.AuthParams) {
			credential, err := azure.NewServicePrincipalCredential(config.AuthParams)
			if err != nil {
				return err
			}
			meta.authContext.credential = credential
			meta.authContext.authType = azurePipelinesAuthTypeServicePrincipal
			return nil
		}
		return fmt.Errorf("no personalAccessToken, Azure service principal, or PodIdentity provider configured")
	case kedav1alpha1.PodIdentityProviderAzureWorkload:
		credential, err := azure.NewChainedCredential(logger, config.PodIdentity)
		if err != nil {
			return err
		}
		meta.authContext.credential = credential
		meta.authContext.authType = azurePipelinesAuthTypeWorkloadIdentity
		return nil
	default:
		return fmt.Errorf("pod identity %s not supported for azure pipelines", config.PodIdentity.Provider)
	}
}

func parseAzurePipelinesMetadata(ctx context.Context, logger logr.Logger, config *scalersconfig.ScalerConfig, httpClient *http.Client) (*azurePipelinesMetadata, error) {
	if config.TriggerMetadata["jobsToFetch"] != "" && config.TriggerMetadata["fetchUnfinishedJobsOnly"] != "" {
		return nil, fmt.Errorf("cannot specify both jobsToFetch and fetchUnfinishedJobsOnly at the same time")
	}
	if config.TriggerMetadata["jobsToFetch"] != "" && config.TriggerMetadata["parent"] != "" {
		return nil, fmt.Errorf("cannot specify both jobsToFetch and parent at the same time")
	}

	meta := &azurePipelinesMetadata{}
	if err := config.TypedConfig(meta); err != nil {
		return nil, fmt.Errorf("error parsing azure pipeline metadata: %w", err)
	}

	meta.triggerIndex = config.TriggerIndex

	if err := configureAzurePipelinesAuth(logger, config, meta); err != nil {
		return nil, err
	}

	if meta.PoolName != "" {
		poolID, err := getPoolIDFromName(ctx, logger, meta.PoolName, meta, httpClient)
		if err != nil {
			return nil, err
		}
		meta.PoolID = poolID
	} else if meta.PoolID != 0 {
		_, err := validatePoolID(ctx, logger, meta.PoolID, meta, httpClient)
		if err != nil {
			return nil, err
		}
	} else {
		return nil, fmt.Errorf("no poolName or poolID given")
	}

	return meta, nil
}

func getPoolIDFromName(ctx context.Context, logger logr.Logger, poolName string, metadata *azurePipelinesMetadata, httpClient *http.Client) (int, error) {
	urlString := fmt.Sprintf("%s/_apis/distributedtask/pools?poolName=%s", metadata.OrganizationURL, url.QueryEscape(poolName))
	body, err := getAzurePipelineRequest(ctx, logger, urlString, metadata, httpClient)

	if err != nil {
		return -1, err
	}

	var result azurePipelinesPoolNameResponse
	err = json.Unmarshal(body, &result)
	if err != nil {
		return -1, err
	}

	count := len(result.Value)
	if count == 0 {
		return -1, fmt.Errorf("agent pool with name `%s` not found", poolName)
	}

	if count != 1 {
		return -1, fmt.Errorf("found %d agent pool with name `%s`", count, poolName)
	}

	return result.Value[0].ID, nil
}

func validatePoolID(ctx context.Context, logger logr.Logger, poolID int, metadata *azurePipelinesMetadata, httpClient *http.Client) (int, error) {
	urlString := fmt.Sprintf("%s/_apis/distributedtask/pools?poolID=%d", metadata.OrganizationURL, poolID)
	body, err := getAzurePipelineRequest(ctx, logger, urlString, metadata, httpClient)

	if err != nil {
		return -1, fmt.Errorf("agent pool with id `%d` not found: %w", poolID, err)
	}

	var result azurePipelinesPoolIDResponse
	err = json.Unmarshal(body, &result)
	if err != nil {
		return -1, err
	}

	return result.ID, nil
}

func getToken(ctx context.Context, metadata *azurePipelinesMetadata, scope string) (string, error) {
	if metadata.authContext.token != nil {
		//if token expires after more then minute from now let's reuse
		if metadata.authContext.token.ExpiresOn.After(time.Now().Add(time.Second * 60)) {
			return metadata.authContext.token.Token, nil
		}
	}
	token, err := metadata.authContext.credential.GetToken(ctx, policy.TokenRequestOptions{
		Scopes: []string{
			scope,
		},
	})

	if err != nil {
		return "", err
	}

	metadata.authContext.token = &token

	return metadata.authContext.token.Token, nil
}

func getAzurePipelineRequest(ctx context.Context, logger logr.Logger, urlString string, metadata *azurePipelinesMetadata, httpClient *http.Client) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", urlString, nil)
	if err != nil {
		return []byte{}, err
	}

	switch metadata.authContext.authType {
	case azurePipelinesAuthTypePAT:
		logger.V(1).Info("making request to ADO REST API using PAT")
		req.SetBasicAuth("", metadata.authContext.pat)
	case azurePipelinesAuthTypeServicePrincipal, azurePipelinesAuthTypeWorkloadIdentity:
		authMethod := string(metadata.authContext.authType)
		logger.V(1).Info("making request to ADO REST API using bearer token", "authMethod", authMethod)
		aadToken, err := getToken(ctx, metadata, devopsResource)
		if err != nil {
			return []byte{}, fmt.Errorf("cannot acquire Azure DevOps token using %s: %w", authMethod, err)
		}
		logger.V(1).Info("token acquired setting auth header as 'bearer XXXXXX'")
		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", aadToken))
	default:
		return []byte{}, fmt.Errorf("no authentication method configured for azure pipelines")
	}

	r, err := httpClient.Do(req)
	if err != nil {
		return []byte{}, err
	}

	b, err := io.ReadAll(r.Body)
	if err != nil {
		return []byte{}, err
	}
	r.Body.Close()

	if r.StatusCode < 200 || r.StatusCode > 299 {
		return []byte{}, fmt.Errorf("the Azure DevOps REST API returned error. urlString: %s status: %d response: %s", urlString, r.StatusCode, string(b))
	}

	// Log when API Rate Limits are reached
	rateLimitRemaining := r.Header[http.CanonicalHeaderKey("X-RateLimit-Remaining")]
	if rateLimitRemaining != nil {
		logger.V(1).Info(fmt.Sprintf("Warning: ADO TSTUs Left %s. When reaching zero requests are delayed, lower the polling interval. See https://learn.microsoft.com/en-us/azure/devops/integrate/concepts/rate-limits?view=azure-devops", rateLimitRemaining))
	}
	rateLimitDelay := r.Header[http.CanonicalHeaderKey("X-RateLimit-Delay")]
	if rateLimitDelay != nil {
		logger.V(1).Info(fmt.Sprintf("Warning: Request to ADO API is delayed by %s seconds. Sending additional requests will increase delay until results are being blocked entirely", rateLimitDelay))
	}

	return b, nil
}

func (s *azurePipelinesScaler) GetAzurePipelinesQueueURL() (string, error) {
	var urlString string
	if s.metadata.FetchUnfinishedJobsOnly {
		// Because completedRequestCount=0 does not change the format of the returned JSON like $top does, we can also use this URL when a `Parent` agent is given
		// However, in order to not force existing users of the `Parent` property to always send the `completedRequestCount=0` query parameter (which may change over time
		//  due it it being an undocumented API), we only send the query parameter when `FetchUnfinishedJobsOnly` is explicitly set to true.
		urlString = fmt.Sprintf("%s/_apis/distributedtask/pools/%d/jobrequests?completedRequestCount=0", s.metadata.OrganizationURL, s.metadata.PoolID)
	} else if s.metadata.Parent != "" {
		// HotFix Issue (#4387), $top changes the format of the returned JSON
		urlString = fmt.Sprintf("%s/_apis/distributedtask/pools/%d/jobrequests", s.metadata.OrganizationURL, s.metadata.PoolID)
	} else {
		urlString = fmt.Sprintf("%s/_apis/distributedtask/pools/%d/jobrequests?$top=%d", s.metadata.OrganizationURL, s.metadata.PoolID, s.metadata.JobsToFetch)
	}

	return urlString, nil
}

func (s *azurePipelinesScaler) GetAzurePipelinesQueueLength(ctx context.Context) (int64, error) {
	urlString, err := s.GetAzurePipelinesQueueURL()
	if err != nil {
		return -1, err
	}

	body, err := getAzurePipelineRequest(ctx, s.logger, urlString, s.metadata, s.httpClient)
	if err != nil {
		return -1, err
	}

	var jrs JobRequests
	err = json.Unmarshal(body, &jrs)
	if err != nil {
		s.logger.Error(err, "Cannot unmarshal ADO JobRequests API response")
		return -1, err
	}

	// for each job check if its parent fulfilled, then demand fulfilled, then finally pool fulfilled
	var count int64
	for _, job := range stripDeadJobs(jrs.Value, s.metadata.ScaleOnInFlight) {
		if s.metadata.Parent == "" && s.metadata.Demands == "" {
			// no plan defined, just add a count
			count++
		} else {
			if s.metadata.Parent == "" {
				// doesn't use parent, switch to demand
				if getCanAgentDemandFulfilJob(job, s.metadata) {
					count++
				}
			} else {
				// does use parent
				if getCanAgentParentFulfilJob(job, s.metadata) {
					count++
				}
			}
		}
	}

	return count, err
}

// stripDeadJobs filters out jobs that should no longer count towards the queue length.
// A job is always skipped when it has finished (Result is set).
// A job that has already been picked up by an agent (ReceiveTime is set) is only
// skipped when scaleOnInFlight is false.
func stripDeadJobs(jobs []JobRequest, scaleOnInFlight bool) []JobRequest {
	var filtered []JobRequest
	for _, job := range jobs {
		if job.Result != nil {
			continue
		}
		if !scaleOnInFlight && !job.ReceiveTime.IsZero() {
			continue
		}
		filtered = append(filtered, job)
	}
	return filtered
}

func stripAgentVFromArray(array []string) []string {
	var result []string

	for _, item := range array {
		if !strings.HasPrefix(item, "Agent.Version") {
			result = append(result, item)
		}
	}
	return result
}

// Determine if the scaledjob has the right demands to spin up
func getCanAgentDemandFulfilJob(jr JobRequest, metadata *azurePipelinesMetadata) bool {
	countDemands := 0
	demandsInJob := stripAgentVFromArray(jr.Demands)
	demandsInScaler := stripAgentVFromArray(strings.Split(metadata.Demands, ","))

	for _, demandInJob := range demandsInJob {
		for _, demandInScaler := range demandsInScaler {
			var isMatch bool
			if metadata.CaseInsensitiveDemandsProcessing {
				isMatch = strings.EqualFold(demandInJob, demandInScaler)
			} else {
				isMatch = demandInJob == demandInScaler
			}
			if isMatch {
				countDemands++
			}
		}
	}

	if metadata.RequireAllDemands {
		return countDemands == len(demandsInJob) && countDemands == len(demandsInScaler)
	} else if metadata.RequireAllDemandsAndIgnoreOthers {
		return countDemands == len(demandsInScaler)
	}
	return countDemands == len(demandsInJob)
}

// Determine if the Job and Parent Agent Template have matching capabilities
func getCanAgentParentFulfilJob(jr JobRequest, metadata *azurePipelinesMetadata) bool {
	matchedAgents := jr.MatchedAgents

	if matchedAgents == nil {
		return false
	}

	for _, m := range *matchedAgents {
		if metadata.Parent == m.Name {
			return true
		}
	}
	return false
}

func (s *azurePipelinesScaler) GetMetricSpecForScaling(context.Context) []v2.MetricSpec {
	externalMetric := &v2.ExternalMetricSource{
		Metric: v2.MetricIdentifier{
			Name: GenerateMetricNameWithIndex(s.metadata.triggerIndex, kedautil.NormalizeString(fmt.Sprintf("azure-pipelines-%d", s.metadata.PoolID))),
		},
		Target: GetMetricTarget(s.metricType, s.metadata.TargetPipelinesQueueLength),
	}
	metricSpec := v2.MetricSpec{External: externalMetric, Type: externalMetricType}
	return []v2.MetricSpec{metricSpec}
}

func (s *azurePipelinesScaler) GetMetricsAndActivity(ctx context.Context, metricName string) ([]external_metrics.ExternalMetricValue, bool, error) {
	queueLen, err := s.GetAzurePipelinesQueueLength(ctx)

	if err != nil {
		s.logger.Error(err, "error getting pipelines queue length")
		return []external_metrics.ExternalMetricValue{}, false, err
	}

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

	return []external_metrics.ExternalMetricValue{metric}, queueLen > s.metadata.ActivationTargetPipelinesQueueLength, nil
}

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