Add a Pusher.PushWithTimeout method, resetting the context timeout

prometheus/push/push.go

@@ -44,6 +44,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"time"
 
 	"github.com/prometheus/common/expfmt"
 	"github.com/prometheus/common/model"
@@ -133,6 +134,19 @@ func (p *Pusher) PushContext(ctx context.Context) error {
 	return p.push(ctx, http.MethodPut)
 }
 
+// PushWithTimeout is like PushContext but overrides the the provided
+// context with a separate timeout. When a job fails due to a
+// higher-level context being canceled (for example, a global timeout)
+// it is easy to reuse that context to push metrics, which fails
+// immediately and metrics pertaining to the original failure will be
+// lost. Use this method for an additional grace period beyond the
+// original context's deadline.
+func (p *Pusher) PushWithTimeout(ctx context.Context, grace time.Duration) error {
+	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), grace)
+	defer cancel()
+	return p.PushContext(ctx)
+}
+
 // Add works like push, but only previously pushed metrics with the same name
 // (and the same job and other grouping labels) will be replaced. (It uses HTTP
 // method “POST” to push to the Pushgateway.)

prometheus/push/push_test.go

@@ -15,11 +15,13 @@ package push
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"io"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/prometheus/common/expfmt"
 
@@ -306,4 +308,16 @@ func TestPush(t *testing.T) {
 	if lastHeader == nil || lastHeader.Get("Authorization") == "" {
 		t.Error("empty Authorization header")
 	}
+
+	// Make sure a canceled context fails, and a reset cancelation
+	// timeout works.
+	canceled, cancel := context.WithCancel(t.Context())
+	cancel()
+	if err := New(pgwOK.URL, "test-timeout").PushContext(canceled); !errors.Is(err, context.Canceled) {
+		t.Errorf("expected canceled push to have failed with %v, not %v", context.Canceled, err)
+	}
+
+	if err := New(pgwOK.URL, "test-timeout").PushWithTimeout(canceled, time.Hour); err != nil {
+		t.Errorf("expected uncanceled push to have succeeded, not %v", err)
+	}
 }