From dc3c14bcdbd4b702f9539738eee205c836240388 Mon Sep 17 00:00:00 2001
From: Ben McNicholl <benmcnicholl@gmail.com>
Date: Fri, 30 Jun 2023 11:42:40 +1000
Subject: [PATCH 1/2] Make timeout a configurable value.

---
 collector/collector.go | 3 ++-
 main.go                | 2 ++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/collector/collector.go b/collector/collector.go
index d24d685e..1b29f3d5 100644
--- a/collector/collector.go
+++ b/collector/collector.go
@@ -34,6 +34,7 @@ type Collector struct {
 	Quiet     bool
 	Debug     bool
 	DebugHttp bool
+	Timeout   int
 }
 
 type Result struct {
@@ -89,7 +90,7 @@ func (c *Collector) Collect() (*Result, error) {
 	}
 
 	httpClient := &http.Client{
-		Timeout: 15 * time.Second,
+		Timeout: time.Duration(c.Timeout) * time.Second,
 	}
 
 	if len(c.Queues) == 0 {
diff --git a/main.go b/main.go
index e5c348de..c2e5b8d7 100644
--- a/main.go
+++ b/main.go
@@ -26,6 +26,7 @@ func main() {
 		debugHttp   = flag.Bool("debug-http", false, "Show full http traces")
 		dryRun      = flag.Bool("dry-run", false, "Whether to only print metrics")
 		endpoint    = flag.String("endpoint", "https://agent.buildkite.com/v3", "A custom Buildkite Agent API endpoint")
+		timeout     = flag.Int("timeout", 15, "Timeout, in seconds, for HTTP requests to Buildkite API")
 
 		// backend config
 		backendOpt     = flag.String("backend", "cloudwatch", "Specify the backend to use: cloudwatch, statsd, prometheus, stackdriver")
@@ -117,6 +118,7 @@ func main() {
 		Quiet:     *quiet,
 		Debug:     *debug,
 		DebugHttp: *debugHttp,
+		Timeout:   *timeout,
 	}
 
 	f := func() (time.Duration, error) {

From 85f26bb709b2d88651921cb95a315f4df1a2a0a3 Mon Sep 17 00:00:00 2001
From: Ben McNicholl <benmcnicholl@gmail.com>
Date: Mon, 3 Jul 2023 15:39:02 +1000
Subject: [PATCH 2/2] Make changes in lambda handler

---
 lambda/main.go | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/lambda/main.go b/lambda/main.go
index 5b19e12a..a2ea4cb7 100644
--- a/lambda/main.go
+++ b/lambda/main.go
@@ -7,6 +7,7 @@ import (
 	"io/ioutil"
 	"log"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 
@@ -56,6 +57,7 @@ func Handler(ctx context.Context, evt json.RawMessage) (string, error) {
 	clwDimensions := os.Getenv("BUILDKITE_CLOUDWATCH_DIMENSIONS")
 	quietString := os.Getenv("BUILDKITE_QUIET")
 	quiet := quietString == "1" || quietString == "true"
+	timeout := os.Getenv("BUILDKITE_AGENT_METRICS_TIMEOUT")
 
 	if quiet {
 		log.SetOutput(ioutil.Discard)
@@ -84,6 +86,16 @@ func Handler(ctx context.Context, evt json.RawMessage) (string, error) {
 		queues = strings.Split(queue, ",")
 	}
 
+	if timeout == "" {
+		timeout = "15"
+	}
+
+	configuredTimeout, err := strconv.Atoi(timeout)
+
+	if err != nil {
+		return "", err
+	}
+
 	userAgent := fmt.Sprintf("buildkite-agent-metrics/%s buildkite-agent-metrics-lambda", version.Version)
 
 	c := collector.Collector{
@@ -94,6 +106,7 @@ func Handler(ctx context.Context, evt json.RawMessage) (string, error) {
 		Quiet:     quiet,
 		Debug:     false,
 		DebugHttp: false,
+		Timeout:   configuredTimeout,
 	}
 
 	switch strings.ToLower(backendOpt) {