Support multiple queue params

README.md

@@ -32,6 +32,12 @@ Restrict it to a single queue with `-queue` if you're scaling a single cluster o
 buildkite-agent-metrics -token abc123 -interval 30s -queue my-queue
 ```
 
+Restrict it to a multiple queues with `-queue`:
+
+```
+buildkite-agent-metrics -token abc123 -interval 30s -queue my-queue1 -queue my-queue2
+```
+
 ### Running as an AWS Lambda
 
 An AWS Lambda bundle is created and published as part of the build process. The lambda will require the [`cloudwatch:PutMetricData`](https://docs.aws.amazon.com/AmazonCloudWatch/latest/DeveloperGuide/publishingMetrics.html) IAM permission.

collector/collector.go

@@ -30,7 +30,7 @@ type Collector struct {
 	Endpoint  string
 	Token     string
 	UserAgent string
-	Queue     string
+	Queues    []string
 	Quiet     bool
 	Debug     bool
 	DebugHttp bool
@@ -88,7 +88,7 @@ func (c *Collector) Collect() (*Result, error) {
 		Queues: map[string]map[string]int{},
 	}
 
-	if c.Queue == "" {
+	if len(c.Queues) == 0 {
 		log.Println("Collecting agent metrics for all queues")
 
 		endpoint, err := url.Parse(c.Endpoint)
@@ -199,64 +199,66 @@ func (c *Collector) Collect() (*Result, error) {
 			result.Queues[queueName][BusyAgentPercentage] = busyAgentPercentage(queueAgentMetrics)
 		}
 	} else {
-		log.Printf("Collecting agent metrics for queue '%s'", c.Queue)
+		for _, queue := range c.Queues {
+			log.Printf("Collecting agent metrics for queue '%s'", queue)
 
-		endpoint, err := url.Parse(c.Endpoint)
-		if err != nil {
-			return nil, err
-		}
+			endpoint, err := url.Parse(c.Endpoint)
+			if err != nil {
+				return nil, err
+			}
 
-		endpoint.Path += "/metrics/queue"
-		endpoint.RawQuery = url.Values{"name": {c.Queue}}.Encode()
+			endpoint.Path += "/metrics/queue"
+			endpoint.RawQuery = url.Values{"name": {queue}}.Encode()
 
-		req, err := http.NewRequest("GET", endpoint.String(), nil)
-		if err != nil {
-			return nil, err
-		}
+			req, err := http.NewRequest("GET", endpoint.String(), nil)
+			if err != nil {
+				return nil, err
+			}
 
-		req.Header.Set("User-Agent", c.UserAgent)
-		req.Header.Set("Authorization", fmt.Sprintf("Token %s", c.Token))
+			req.Header.Set("User-Agent", c.UserAgent)
+			req.Header.Set("Authorization", fmt.Sprintf("Token %s", c.Token))
 
-		if c.DebugHttp {
-			if dump, err := httputil.DumpRequest(req, true); err == nil {
-				log.Printf("DEBUG request uri=%s\n%s\n", req.URL, dump)
+			if c.DebugHttp {
+				if dump, err := httputil.DumpRequest(req, true); err == nil {
+					log.Printf("DEBUG request uri=%s\n%s\n", req.URL, dump)
+				}
 			}
-		}
 
-		res, err := http.DefaultClient.Do(req)
-		if err != nil {
-			return nil, err
-		}
+			res, err := http.DefaultClient.Do(req)
+			if err != nil {
+				return nil, err
+			}
 
-		if c.DebugHttp {
-			if dump, err := httputil.DumpResponse(res, true); err == nil {
-				log.Printf("DEBUG response uri=%s\n%s\n", req.URL, dump)
+			if c.DebugHttp {
+				if dump, err := httputil.DumpResponse(res, true); err == nil {
+					log.Printf("DEBUG response uri=%s\n%s\n", req.URL, dump)
+				}
 			}
-		}
 
-		var queueMetrics queueMetricsResponse
-		defer res.Body.Close()
-		err = json.NewDecoder(res.Body).Decode(&queueMetrics)
-		if err != nil {
-			return nil, err
-		}
+			var queueMetrics queueMetricsResponse
+			defer res.Body.Close()
+			err = json.NewDecoder(res.Body).Decode(&queueMetrics)
+			if err != nil {
+				return nil, err
+			}
 
-		if queueMetrics.Organization.Slug == "" {
-			return nil, fmt.Errorf("No organization slug was found in the metrics response")
-		}
+			if queueMetrics.Organization.Slug == "" {
+				return nil, fmt.Errorf("No organization slug was found in the metrics response")
+			}
 
-		log.Printf("Found organization %q", queueMetrics.Organization.Slug)
-		result.Org = queueMetrics.Organization.Slug
-
-		result.Queues[c.Queue] = map[string]int{
-			ScheduledJobsCount:  queueMetrics.Jobs.Scheduled,
-			RunningJobsCount:    queueMetrics.Jobs.Running,
-			UnfinishedJobsCount: queueMetrics.Jobs.Total,
-			WaitingJobsCount:    queueMetrics.Jobs.Waiting,
-			IdleAgentCount:      queueMetrics.Agents.Idle,
-			BusyAgentCount:      queueMetrics.Agents.Busy,
-			TotalAgentCount:     queueMetrics.Agents.Total,
-			BusyAgentPercentage: busyAgentPercentage(queueMetrics.Agents),
+			log.Printf("Found organization %q", queueMetrics.Organization.Slug)
+			result.Org = queueMetrics.Organization.Slug
+
+			result.Queues[queue] = map[string]int{
+				ScheduledJobsCount:  queueMetrics.Jobs.Scheduled,
+				RunningJobsCount:    queueMetrics.Jobs.Running,
+				UnfinishedJobsCount: queueMetrics.Jobs.Total,
+				WaitingJobsCount:    queueMetrics.Jobs.Waiting,
+				IdleAgentCount:      queueMetrics.Agents.Idle,
+				BusyAgentCount:      queueMetrics.Agents.Busy,
+				TotalAgentCount:     queueMetrics.Agents.Total,
+				BusyAgentPercentage: busyAgentPercentage(queueMetrics.Agents),
+			}
 		}
 	}
 

collector/collector_test.go

@@ -247,7 +247,7 @@ func TestCollectorWithSomeJobsAndAgentsForAQueue(t *testing.T) {
 		Endpoint:  s.URL,
 		Token:     "abc123",
 		UserAgent: "some-client/1.2.3",
-		Queue:     "deploy",
+		Queues:    []string{"deploy"},
 	}
 	res, err := c.Collect()
 	if err != nil {

lambda/main.go

@@ -63,7 +63,7 @@ func Handler(ctx context.Context, evt json.RawMessage) (string, error) {
 		UserAgent: userAgent,
 		Endpoint:  "https://agent.buildkite.com/v3",
 		Token:     token,
-		Queue:     queue,
+		Queues:    []string{queue},
 		Quiet:     quiet,
 		Debug:     false,
 		DebugHttp: false,

main.go

@@ -38,11 +38,12 @@ func main() {
 		gcpProjectID   = flag.String("stackdriver-projectid", "", "Specify Stackdriver Project ID")
 		nrAppName      = flag.String("newrelic-app-name", "", "New Relic application name for metric events")
 		nrLicenseKey   = flag.String("newrelic-license-key", "", "New Relic license key for publishing events")
-
-		// filters
-		queue = flag.String("queue", "", "Only include a specific queue")
 	)
 
+	// custom config for multiple queues
+	var queues stringSliceFlag
+	flag.Var(&queues, "queue", "Specific queues to process")
+
 	flag.Parse()
 
 	if *showVersion {
@@ -108,7 +109,7 @@ func main() {
 		UserAgent: userAgent,
 		Endpoint:  *endpoint,
 		Token:     *token,
-		Queue:     *queue,
+		Queues:    []string(queues),
 		Quiet:     *quiet,
 		Debug:     *debug,
 		DebugHttp: *debugHttp,
@@ -160,3 +161,14 @@ func main() {
 		}
 	}
 }
+
+type stringSliceFlag []string
+
+func (i *stringSliceFlag) String() string {
+	return fmt.Sprintf("%v", *i)
+}
+
+func (i *stringSliceFlag) Set(value string) error {
+	*i = append(*i, value)
+	return nil
+}