<commit_before>package main

import (
	"flag"
	"log"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	buildkite "gopkg.in/buildkite/go-buildkite.v2"
)

var queuePattern *regexp.Regexp

func init() {
	queuePattern = regexp.MustCompile(`(?i)^queue=(.+?)$`)
}

// Generates:

// Buildkite > RunningBuildsCount
// Buildkite > RunningJobsCount
// Buildkite > ScheduledBuildsCount
// Buildkite > ScheduledJobsCount

// Buildkite > (Queue) > RunningBuildsCount
// Buildkite > (Queue) > RunningJobsCount
// Buildkite > (Queue) > ScheduledBuildsCount
// Buildkite > (Queue) > ScheduledJobsCount

// Buildkite > (Pipeline) > RunningBuildsCount
// Buildkite > (Pipeline) > RunningJobsCount
// Buildkite > (Pipeline) > ScheduledBuildsCount
// Buildkite > (Pipeline) > ScheduledJobsCount

func main() {
	var (
		accessToken = flag.String("token", "", "A Buildkite API Access Token")
		orgSlug     = flag.String("org", "", "A Buildkite Organization Slug")
		interval    = flag.Duration("interval", 0, "Update metrics every interval, rather than once")
		debug       = flag.Bool("debug", false, "Show API debugging output")
	)

	flag.Parse()

	if *accessToken == "" {
		log.Fatal("Must provide a value for -token")
	}

	if *orgSlug == "" {
		log.Fatal("Must provide a value for -org")
	}

	config, err := buildkite.NewTokenConfig(*accessToken, false)
	if err != nil {
		log.Fatalf("client config failed: %s", err)
	}

	client := buildkite.NewClient(config.Client())
	buildkite.SetHttpDebug(*debug)

	if err := runCollector(client, *orgSlug, time.Hour*24); err != nil {
		log.Fatal(err)
	}

	if *interval > 0 {
		for _ = range time.NewTicker(*interval).C {
			if err := runCollector(client, *orgSlug, time.Hour); err != nil {
				log.Println(err)
			}
		}
	}
}

func runCollector(client *buildkite.Client, orgSlug string, historical time.Duration) error {
	svc := cloudwatch.New(session.New())

	res := &result{
		totals:    newCounts(),
		queues:    map[string]counts{},
		pipelines: map[string]counts{},
	}

	log.Printf("Collecting buildkite metrics from org %s", orgSlug)
	if err := res.addBuildAndJobMetrics(client, orgSlug, historical); err != nil {
		return err
	}

	metrics := res.toMetrics()
	log.Printf("Extracted %d cloudwatch metrics from results", len(metrics))

	for _, chunk := range chunkMetricData(10, metrics) {
		log.Printf("Submitting chunk of %d metrics to Cloudwatch", len(chunk))
		if err := putMetricData(svc, chunk); err != nil {
			if err != nil {
				return err
			}
		}
	}

	return nil
}

const (
	runningBuildsCount   = "RunningBuildsCount"
	runningJobsCount     = "RunningJobsCount"
	scheduledBuildsCount = "ScheduledBuildsCount"
	scheduledJobsCount   = "ScheduledJobsCount"
)

type counts map[string]int

func newCounts() counts {
	return counts{
		runningBuildsCount:   0,
		scheduledBuildsCount: 0,
		runningJobsCount:     0,
		scheduledJobsCount:   0,
	}
}

func (c counts) toMetrics(dimensions []*cloudwatch.Dimension) []*cloudwatch.MetricDatum {
	m := []*cloudwatch.MetricDatum{}

	for k, v := range c {
		m = append(m, &cloudwatch.MetricDatum{
			MetricName: aws.String(k),
			Dimensions: dimensions,
			Value:      aws.Float64(float64(v)),
			Unit:       aws.String("Count"),
		})
	}

	return m
}

func queue(j *buildkite.Job) string {
	for _, m := range j.AgentQueryRules {
		if match := queuePattern.FindStringSubmatch(m); match != nil {
			return match[1]
		}
	}
	return "default"
}

func uniqueQueues(builds []buildkite.Build) []string {
	queueMap := map[string]struct{}{}

	for _, b := range builds {
		for _, j := range b.Jobs {
			queueMap[queue(j)] = struct{}{}
		}
	}

	queues := []string{}

	for q := range queueMap {
		queues = append(queues, q)
	}

	return queues
}

type result struct {
	totals            counts
	queues, pipelines map[string]counts
}

func (r *result) toMetrics() []*cloudwatch.MetricDatum {
	data := []*cloudwatch.MetricDatum{}
	data = append(data, r.totals.toMetrics(nil)...)

	for name, c := range r.queues {
		data = append(data, c.toMetrics([]*cloudwatch.Dimension{
			{Name: aws.String("Queue"), Value: aws.String(name)},
		})...)
	}

	for name, c := range r.pipelines {
		data = append(data, c.toMetrics([]*cloudwatch.Dimension{
			{Name: aws.String("Pipeline"), Value: aws.String(name)},
		})...)
	}

	return data
}

func (r *result) addBuildAndJobMetrics(client *buildkite.Client, orgSlug string, historical time.Duration) error {
	// Algorithm:
	// Get Builds with finished_from = 24 hours ago
	// Build results with zero values for pipelines/queues
	// Get all running and scheduled builds, add to results

	// TODO: page this
	builds, _, err := client.Builds.ListByOrg(orgSlug, &buildkite.BuildsListOptions{
		FinishedFrom: time.Now().UTC().Add(historical * -1),
	})
	if err != nil {
		return err
	}

	for _, queue := range uniqueQueues(builds) {
		r.queues[queue] = newCounts()
	}

	for _, build := range builds {
		r.pipelines[*build.Pipeline.Name] = newCounts()
	}

	p := &pager{
		lister: func(page int) (interface{}, int, error) {
			builds, resp, err := client.Builds.ListByOrg(orgSlug, &buildkite.BuildsListOptions{
				State: []string{"scheduled", "running"},
				ListOptions: buildkite.ListOptions{
					Page: page,
				},
			})
			log.Printf("Builds page %d has %d builds, next page is %d",
				page,
				len(builds),
				resp.NextPage,
			)
			return builds, resp.NextPage, err
		},
	}

	return p.Pages(func(v interface{}, lastPage bool) bool {
		for _, build := range v.([]buildkite.Build) {
			log.Printf("Adding build to stats (id=%q, pipeline=%q, branch=%q, state=%q)",
				*build.ID, *build.Pipeline.Name, *build.Branch, *build.State)

			if _, ok := r.pipelines[*build.Pipeline.Name]; !ok {
				r.pipelines[*build.Pipeline.Name] = newCounts()
			}

			switch *build.State {
			case "running":
				r.totals[runningBuildsCount]++
				r.pipelines[*build.Pipeline.Name][runningBuildsCount]++

			case "scheduled":
				r.totals[scheduledBuildsCount]++
				r.pipelines[*build.Pipeline.Name][scheduledBuildsCount]++
			}

			var buildQueues = map[string]int{}

			for _, job := range build.Jobs {
				if job.Type != nil && *job.Type == "waiter" {
					continue
				}

				state := ""
				if job.State != nil {
					state = *job.State
				}

				log.Printf("Adding job to stats (id=%q, pipeline=%q, queue=%q, type=%q, state=%q)",
					*job.ID, *build.Pipeline.Name, queue(job), *job.Type, state)

				if _, ok := r.queues[queue(job)]; !ok {
					r.queues[queue(job)] = newCounts()
				}

				switch state {
				case "running":
					r.totals[runningJobsCount]++
					r.queues[queue(job)][runningJobsCount]++

				case "scheduled":
					r.totals[scheduledJobsCount]++
					r.queues[queue(job)][scheduledJobsCount]++
				}

				buildQueues[queue(job)]++
			}

			// add build metrics to queues
			if len(buildQueues) > 0 {
				for queue := range buildQueues {
					switch *build.State {
					case "running":
						r.queues[queue][runningBuildsCount]++

					case "scheduled":
						r.queues[queue][scheduledBuildsCount]++
					}
				}
			}
		}
		return true
	})
}

type pager struct {
	lister func(page int) (v interface{}, nextPage int, err error)
}

func (p *pager) Pages(f func(v interface{}, lastPage bool) bool) error {
	page := 1
	for {
		val, nextPage, err := p.lister(page)
		if err != nil {
			return err
		}
		if !f(val, nextPage == 0) || nextPage == 0 {
			break
		}
		page = nextPage
	}
	return nil
}

func chunkMetricData(size int, data []*cloudwatch.MetricDatum) [][]*cloudwatch.MetricDatum {
	var chunks = [][]*cloudwatch.MetricDatum{}
	for i := 0; i < len(data); i += size {
		end := i + size
		if end > len(data) {
			end = len(data)
		}
		chunks = append(chunks, data[i:end])
	}
	return chunks
}

func putMetricData(svc *cloudwatch.CloudWatch, data []*cloudwatch.MetricDatum) error {
	_, err := svc.PutMetricData(&cloudwatch.PutMetricDataInput{
		MetricData: data,
		Namespace:  aws.String("Buildkite"),
	})
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Use correct import path for go-buildkite<commit_after>package main

import (
	"flag"
	"log"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"gopkg.in/buildkite/go-buildkite.v2/buildkite"
)

var queuePattern *regexp.Regexp

func init() {
	queuePattern = regexp.MustCompile(`(?i)^queue=(.+?)$`)
}

// Generates:

// Buildkite > RunningBuildsCount
// Buildkite > RunningJobsCount
// Buildkite > ScheduledBuildsCount
// Buildkite > ScheduledJobsCount

// Buildkite > (Queue) > RunningBuildsCount
// Buildkite > (Queue) > RunningJobsCount
// Buildkite > (Queue) > ScheduledBuildsCount
// Buildkite > (Queue) > ScheduledJobsCount

// Buildkite > (Pipeline) > RunningBuildsCount
// Buildkite > (Pipeline) > RunningJobsCount
// Buildkite > (Pipeline) > ScheduledBuildsCount
// Buildkite > (Pipeline) > ScheduledJobsCount

func main() {
	var (
		accessToken = flag.String("token", "", "A Buildkite API Access Token")
		orgSlug     = flag.String("org", "", "A Buildkite Organization Slug")
		interval    = flag.Duration("interval", 0, "Update metrics every interval, rather than once")
		debug       = flag.Bool("debug", false, "Show API debugging output")
	)

	flag.Parse()

	if *accessToken == "" {
		log.Fatal("Must provide a value for -token")
	}

	if *orgSlug == "" {
		log.Fatal("Must provide a value for -org")
	}

	config, err := buildkite.NewTokenConfig(*accessToken, false)
	if err != nil {
		log.Fatalf("client config failed: %s", err)
	}

	client := buildkite.NewClient(config.Client())
	buildkite.SetHttpDebug(*debug)

	if err := runCollector(client, *orgSlug, time.Hour*24); err != nil {
		log.Fatal(err)
	}

	if *interval > 0 {
		for _ = range time.NewTicker(*interval).C {
			if err := runCollector(client, *orgSlug, time.Hour); err != nil {
				log.Println(err)
			}
		}
	}
}

func runCollector(client *buildkite.Client, orgSlug string, historical time.Duration) error {
	svc := cloudwatch.New(session.New())

	res := &result{
		totals:    newCounts(),
		queues:    map[string]counts{},
		pipelines: map[string]counts{},
	}

	log.Printf("Collecting buildkite metrics from org %s", orgSlug)
	if err := res.addBuildAndJobMetrics(client, orgSlug, historical); err != nil {
		return err
	}

	metrics := res.toMetrics()
	log.Printf("Extracted %d cloudwatch metrics from results", len(metrics))

	for _, chunk := range chunkMetricData(10, metrics) {
		log.Printf("Submitting chunk of %d metrics to Cloudwatch", len(chunk))
		if err := putMetricData(svc, chunk); err != nil {
			if err != nil {
				return err
			}
		}
	}

	return nil
}

const (
	runningBuildsCount   = "RunningBuildsCount"
	runningJobsCount     = "RunningJobsCount"
	scheduledBuildsCount = "ScheduledBuildsCount"
	scheduledJobsCount   = "ScheduledJobsCount"
)

type counts map[string]int

func newCounts() counts {
	return counts{
		runningBuildsCount:   0,
		scheduledBuildsCount: 0,
		runningJobsCount:     0,
		scheduledJobsCount:   0,
	}
}

func (c counts) toMetrics(dimensions []*cloudwatch.Dimension) []*cloudwatch.MetricDatum {
	m := []*cloudwatch.MetricDatum{}

	for k, v := range c {
		m = append(m, &cloudwatch.MetricDatum{
			MetricName: aws.String(k),
			Dimensions: dimensions,
			Value:      aws.Float64(float64(v)),
			Unit:       aws.String("Count"),
		})
	}

	return m
}

func queue(j *buildkite.Job) string {
	for _, m := range j.AgentQueryRules {
		if match := queuePattern.FindStringSubmatch(m); match != nil {
			return match[1]
		}
	}
	return "default"
}

func uniqueQueues(builds []buildkite.Build) []string {
	queueMap := map[string]struct{}{}

	for _, b := range builds {
		for _, j := range b.Jobs {
			queueMap[queue(j)] = struct{}{}
		}
	}

	queues := []string{}

	for q := range queueMap {
		queues = append(queues, q)
	}

	return queues
}

type result struct {
	totals            counts
	queues, pipelines map[string]counts
}

func (r *result) toMetrics() []*cloudwatch.MetricDatum {
	data := []*cloudwatch.MetricDatum{}
	data = append(data, r.totals.toMetrics(nil)...)

	for name, c := range r.queues {
		data = append(data, c.toMetrics([]*cloudwatch.Dimension{
			{Name: aws.String("Queue"), Value: aws.String(name)},
		})...)
	}

	for name, c := range r.pipelines {
		data = append(data, c.toMetrics([]*cloudwatch.Dimension{
			{Name: aws.String("Pipeline"), Value: aws.String(name)},
		})...)
	}

	return data
}

func (r *result) addBuildAndJobMetrics(client *buildkite.Client, orgSlug string, historical time.Duration) error {
	// Algorithm:
	// Get Builds with finished_from = 24 hours ago
	// Build results with zero values for pipelines/queues
	// Get all running and scheduled builds, add to results

	// TODO: page this
	builds, _, err := client.Builds.ListByOrg(orgSlug, &buildkite.BuildsListOptions{
		FinishedFrom: time.Now().UTC().Add(historical * -1),
	})
	if err != nil {
		return err
	}

	for _, queue := range uniqueQueues(builds) {
		r.queues[queue] = newCounts()
	}

	for _, build := range builds {
		r.pipelines[*build.Pipeline.Name] = newCounts()
	}

	p := &pager{
		lister: func(page int) (interface{}, int, error) {
			builds, resp, err := client.Builds.ListByOrg(orgSlug, &buildkite.BuildsListOptions{
				State: []string{"scheduled", "running"},
				ListOptions: buildkite.ListOptions{
					Page: page,
				},
			})
			log.Printf("Builds page %d has %d builds, next page is %d",
				page,
				len(builds),
				resp.NextPage,
			)
			return builds, resp.NextPage, err
		},
	}

	return p.Pages(func(v interface{}, lastPage bool) bool {
		for _, build := range v.([]buildkite.Build) {
			log.Printf("Adding build to stats (id=%q, pipeline=%q, branch=%q, state=%q)",
				*build.ID, *build.Pipeline.Name, *build.Branch, *build.State)

			if _, ok := r.pipelines[*build.Pipeline.Name]; !ok {
				r.pipelines[*build.Pipeline.Name] = newCounts()
			}

			switch *build.State {
			case "running":
				r.totals[runningBuildsCount]++
				r.pipelines[*build.Pipeline.Name][runningBuildsCount]++

			case "scheduled":
				r.totals[scheduledBuildsCount]++
				r.pipelines[*build.Pipeline.Name][scheduledBuildsCount]++
			}

			var buildQueues = map[string]int{}

			for _, job := range build.Jobs {
				if job.Type != nil && *job.Type == "waiter" {
					continue
				}

				state := ""
				if job.State != nil {
					state = *job.State
				}

				log.Printf("Adding job to stats (id=%q, pipeline=%q, queue=%q, type=%q, state=%q)",
					*job.ID, *build.Pipeline.Name, queue(job), *job.Type, state)

				if _, ok := r.queues[queue(job)]; !ok {
					r.queues[queue(job)] = newCounts()
				}

				switch state {
				case "running":
					r.totals[runningJobsCount]++
					r.queues[queue(job)][runningJobsCount]++

				case "scheduled":
					r.totals[scheduledJobsCount]++
					r.queues[queue(job)][scheduledJobsCount]++
				}

				buildQueues[queue(job)]++
			}

			// add build metrics to queues
			if len(buildQueues) > 0 {
				for queue := range buildQueues {
					switch *build.State {
					case "running":
						r.queues[queue][runningBuildsCount]++

					case "scheduled":
						r.queues[queue][scheduledBuildsCount]++
					}
				}
			}
		}
		return true
	})
}

type pager struct {
	lister func(page int) (v interface{}, nextPage int, err error)
}

func (p *pager) Pages(f func(v interface{}, lastPage bool) bool) error {
	page := 1
	for {
		val, nextPage, err := p.lister(page)
		if err != nil {
			return err
		}
		if !f(val, nextPage == 0) || nextPage == 0 {
			break
		}
		page = nextPage
	}
	return nil
}

func chunkMetricData(size int, data []*cloudwatch.MetricDatum) [][]*cloudwatch.MetricDatum {
	var chunks = [][]*cloudwatch.MetricDatum{}
	for i := 0; i < len(data); i += size {
		end := i + size
		if end > len(data) {
			end = len(data)
		}
		chunks = append(chunks, data[i:end])
	}
	return chunks
}

func putMetricData(svc *cloudwatch.CloudWatch, data []*cloudwatch.MetricDatum) error {
	_, err := svc.PutMetricData(&cloudwatch.PutMetricDataInput{
		MetricData: data,
		Namespace:  aws.String("Buildkite"),
	})
	if err != nil {
		return err
	}

	return nil
}

{'fraction_non_alphanumeric': 0.1528062015503876, 'fraction_numerical': 0.0029767441860465114, 'mean_word_length': 3.6166618952190093, 'pattern_counts': {'":': 12, '<': 5, '<?xml version=': 0, '>': 49, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '237805', 'n_tokens_mistral': 6405, 'n_tokens_neox': 5734, 'n_words': 1300}