use contexts properly with timeout

timeouts could cause the code to exit, and started before the action was taken. create contexts with timeouts when necessary, and use them in prometheusClient and golang net.http API calls

Author: Ahmet Oeztuerk

internal/helper/prometheus.go

@@ -1,6 +1,7 @@
 package helper
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"io"
@@ -93,7 +94,7 @@ func NewAPIClientV1(address *url.URL) (v1.API, error) {
 }
 
 // DoAPIRequest does the http handling for an api request
-func DoAPIRequest(url *url.URL) ([]byte, error) {
+func DoAPIRequest(ctx context.Context, url *url.URL) ([]byte, error) {
 	transport := http.DefaultTransport.(*http.Transport).Clone()
 	transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: InsecureSkipVerify}
 
@@ -107,7 +108,13 @@ func DoAPIRequest(url *url.URL) ([]byte, error) {
 		httpClient.Jar.SetCookies(url, Cookies)
 	}
 
-	resp, err := httpClient.Get(url.String())
+	// Create request with context to support timeout
+	req, err := http.NewRequestWithContext(ctx, "GET", url.String(), nil)
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := httpClient.Do(req)
 	if err != nil {
 		return nil, err
 	}

internal/mode/ping.go

@@ -26,7 +26,7 @@ type buildInfo struct {
 }
 
 // Ping will fetch build information from the prometheus server
-func Ping(address *url.URL, collection *check_x.PerformanceDataCollection) (check_x.State, string, error) {
+func Ping(ctx context.Context, address *url.URL, collection *check_x.PerformanceDataCollection) (check_x.State, string, error) {
 	if address == nil {
 		err := fmt.Errorf("address to query is null")
 		return check_x.Unknown, fmt.Sprintf("Error: %s", err.Error()), err
@@ -44,7 +44,7 @@ func Ping(address *url.URL, collection *check_x.PerformanceDataCollection) (chec
 
 	query := `prometheus_build_info{job="prometheus"}`
 	startTime := time.Now()
-	result, _, err := apiClient.Query(context.TODO(), query, time.Now())
+	result, _, err := apiClient.Query(ctx, query, time.Now())
 	endTime := time.Now()
 	if err != nil {
 		return check_x.Unknown, fmt.Sprintf("Error when querying API: %s", err.Error()), err

internal/mode/query.go

@@ -17,7 +17,7 @@ import (
 )
 
 // Query allows the user to test data in the prometheus server
-func Query(address *url.URL, query, warning, critical, alias, search, replace, emptyQueryMessage string, emptyQueryStatus check_x.State, collection *check_x.PerformanceDataCollection) (check_x.State, string, error) {
+func Query(ctx context.Context, address *url.URL, query, warning, critical, alias, search, replace, emptyQueryMessage string, emptyQueryStatus check_x.State, collection *check_x.PerformanceDataCollection) (check_x.State, string, error) {
 	if address == nil {
 		err := fmt.Errorf("address to query is null")
 		return check_x.Unknown, fmt.Sprintf("Error: %s", err.Error()), err
@@ -51,7 +51,7 @@ func Query(address *url.URL, query, warning, critical, alias, search, replace, e
 		return check_x.Unknown, fmt.Sprintf("Error creating apiClient: %s", err.Error()), err
 	}
 
-	result, _, err := apiClient.Query(context.TODO(), query, time.Now())
+	result, _, err := apiClient.Query(ctx, query, time.Now())
 	if err != nil {
 		return check_x.Unknown, fmt.Sprintf("Error when querying: %s", err.Error()), err
 	}

internal/mode/targetsHealth.go

@@ -1,6 +1,7 @@
 package mode
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/url"
@@ -36,13 +37,13 @@ type targets struct {
 	} `json:"data"`
 }
 
-func getTargets(address *url.URL) (*targets, error) {
-	u, err := url.Parse(address.String())
+func getTargets(ctx context.Context, address *url.URL) (*targets, error) {
+	url, err := url.Parse(address.String())
 	if err != nil {
 		return nil, err
 	}
-	u.Path = path.Join(u.Path, "/api/v1/targets")
-	jsonBytes, err := helper.DoAPIRequest(u)
+	url.Path = path.Join(url.Path, "/api/v1/targets")
+	jsonBytes, err := helper.DoAPIRequest(ctx, url)
 	if err != nil {
 		return nil, err
 	}
@@ -55,7 +56,7 @@ func getTargets(address *url.URL) (*targets, error) {
 }
 
 // TargetsHealth tests the health of the targets
-func TargetsHealth(address *url.URL, label, warning, critical string, collection *check_x.PerformanceDataCollection) (check_x.State, string, error) {
+func TargetsHealth(ctx context.Context, address *url.URL, label, warning, critical string, collection *check_x.PerformanceDataCollection) (check_x.State, string, error) {
 	if address == nil {
 		err := fmt.Errorf("address to query is null")
 		return check_x.Unknown, fmt.Sprintf("Error: %s", err.Error()), err
@@ -76,7 +77,7 @@ func TargetsHealth(address *url.URL, label, warning, critical string, collection
 		return check_x.Unknown, fmt.Sprintf("Error creating critThreshold from '%s' : %s", critical, err.Error()), err
 	}
 
-	targets, err := getTargets(address)
+	targets, err := getTargets(ctx, address)
 	if err != nil {
 		return check_x.Unknown, fmt.Sprintf("Error getting targets out of address: %s : %s", address.String(), err.Error()), err
 	}

pkg/checker/checker.go

@@ -41,12 +41,6 @@ var (
 	emptyQueryStatus    check_x.State
 )
 
-func startTimeout() {
-	if timeout != 0 {
-		check_x.StartTimeout(time.Duration(timeout) * time.Second)
-	}
-}
-
 // This function is intended to be used for single-use cli mode
 // It will be called from main executable function as it returns int
 func CheckMain(args []string) int {
@@ -71,10 +65,10 @@ func GenerateStdout(state check_x.State, msg string, collection *check_x.Perform
 // It can be used as a library import
 func Check(args []string) (check_x.State, string, *check_x.PerformanceDataCollection, error) {
 
-	var state check_x.State
-	var msg string
+	state := check_x.Unknown
+	msg := "Cli action did not run yet"
 	var collection = check_x.NewPerformanceDataCollection()
-	var err error
+	var err error = nil
 
 	cmd := &cli.Command{
 		Name:    "check_prometheus",
@@ -118,8 +112,15 @@ func Check(args []string) (check_x.State, string, *check_x.PerformanceDataCollec
 						Usage:       "Returns the build informations",
 						Description: `This check requires that the prometheus server itself is listed as target. Following query will be used: 'prometheus_build_info{job="prometheus"}'`,
 						Action: func(ctx context.Context, cmd *cli.Command) error {
-							startTimeout()
-							state, msg, err = mode.Ping(address, &collection)
+							var ctxPing context.Context
+							var ctxPingCancel context.CancelFunc
+							if timeout == 0 {
+								ctxPing = context.WithoutCancel(ctx)
+							} else {
+								ctxPing, ctxPingCancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
+								defer ctxPingCancel()
+							}
+							state, msg, err = mode.Ping(ctxPing, address, &collection)
 							return err
 						},
 						Flags: []cli.Flag{
@@ -193,9 +194,17 @@ func Check(args []string) (check_x.State, string, *check_x.PerformanceDataCollec
 											--> UNKNOWN - The given States do not contain an State
 
 										`,
-						Action: func(c context.Context, cmd *cli.Command) error {
-							startTimeout()
-							state, msg, err = mode.Query(address, queryDecoded, warning, critical, alias, search, replace, emptyQueryMessage, emptyQueryStatus, &collection)
+						Action: func(ctx context.Context, cmd *cli.Command) error {
+							var ctxQuery context.Context
+							var ctxQueryCancel context.CancelFunc
+							if timeout == 0 {
+								ctxQuery = context.WithoutCancel(ctx)
+							} else {
+								ctxQuery, ctxQueryCancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
+								defer ctxQueryCancel()
+							}
+
+							state, msg, err = mode.Query(ctxQuery, address, queryDecoded, warning, critical, alias, search, replace, emptyQueryMessage, emptyQueryStatus, &collection)
 							return err
 						},
 						Flags: []cli.Flag{
@@ -360,9 +369,17 @@ func Check(args []string) (check_x.State, string, *check_x.PerformanceDataCollec
 						HideHelp:    false,
 						Usage:       "Returns the health of the targets",
 						Description: `The warning and critical thresholds are appied on the health_rate. The health_rate is calculted: sum(healthy) / sum(targets).`,
-						Action: func(c context.Context, cmd *cli.Command) error {
-							startTimeout()
-							state, msg, err = mode.TargetsHealth(address, label, warning, critical, &collection)
+						Action: func(ctx context.Context, cmd *cli.Command) error {
+							var ctxTargetsHealth context.Context
+							var ctxTargetsHealthCancel context.CancelFunc
+							if timeout == 0 {
+								ctxTargetsHealth = context.WithoutCancel(ctx)
+							} else {
+								ctxTargetsHealth, ctxTargetsHealthCancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
+								defer ctxTargetsHealthCancel()
+							}
+
+							state, msg, err = mode.TargetsHealth(ctxTargetsHealth, address, label, warning, critical, &collection)
 							return err
 						},
 						Flags: []cli.Flag{