Merge 49d2002e7c into 13dd526f11

Fixes https://github.com/coder/coder/issues/13144

cli/login.go

@@ -287,7 +287,8 @@ func (r *RootCmd) login() *serpent.Command {
 				}
 
 				sessionToken, err = cliui.Prompt(inv, cliui.PromptOptions{
-					Text: "Paste your token here:",
+					Text:   "Paste your token here:",
+					Secret: true,
 					Validate: func(token string) error {
 						client.SetSessionToken(token)
 						_, err := client.User(ctx, codersdk.Me)

cli/server.go

@@ -1441,7 +1441,7 @@ func newProvisionerDaemon(
 
 			connector[string(database.ProvisionerTypeTerraform)] = sdkproto.NewDRPCProvisionerClient(terraformClient)
 		default:
-			return nil, fmt.Errorf("unknown provisioner type %q", provisionerType)
+			return nil, xerrors.Errorf("unknown provisioner type %q", provisionerType)
 		}
 	}
 

cli/ssh.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"net/url"
 	"os"
@@ -79,6 +80,10 @@ func (r *RootCmd) ssh() *serpent.Command {
 			ctx, cancel := context.WithCancel(ctx)
 			defer cancel()
 
+			// Prevent unnecessary logs from the stdlib from messing up the TTY.
+			// See: https://github.com/coder/coder/issues/13144
+			log.SetOutput(io.Discard)
+
 			logger := inv.Logger
 			defer func() {
 				if retErr != nil {

coderd/database/pubsub/latency.go

@@ -0,0 +1,74 @@
+package pubsub
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog"
+)
+
+// LatencyMeasurer is used to measure the send & receive latencies of the underlying Pubsub implementation. We use these
+// measurements to export metrics which can indicate when a Pubsub implementation's queue is overloaded and/or full.
+type LatencyMeasurer struct {
+	// Create unique pubsub channel names so that multiple coderd replicas do not clash when performing latency measurements.
+	channel uuid.UUID
+	logger  slog.Logger
+}
+
+// LatencyMessageLength is the length of a UUIDv4 encoded to hex.
+const LatencyMessageLength = 36
+
+func NewLatencyMeasurer(logger slog.Logger) *LatencyMeasurer {
+	return &LatencyMeasurer{
+		channel: uuid.New(),
+		logger:  logger,
+	}
+}
+
+// Measure takes a given Pubsub implementation, publishes a message & immediately receives it, and returns the observed latency.
+func (lm *LatencyMeasurer) Measure(ctx context.Context, p Pubsub) (send float64, recv float64, err error) {
+	var (
+		start time.Time
+		res   = make(chan float64, 1)
+	)
+
+	msg := []byte(uuid.New().String())
+
+	cancel, err := p.Subscribe(lm.latencyChannelName(), func(ctx context.Context, in []byte) {
+		if !bytes.Equal(in, msg) {
+			lm.logger.Warn(ctx, "received unexpected message", slog.F("got", in), slog.F("expected", msg))
+			return
+		}
+
+		res <- time.Since(start).Seconds()
+	})
+	if err != nil {
+		return -1, -1, xerrors.Errorf("failed to subscribe: %w", err)
+	}
+	defer cancel()
+
+	start = time.Now()
+	err = p.Publish(lm.latencyChannelName(), msg)
+	if err != nil {
+		return -1, -1, xerrors.Errorf("failed to publish: %w", err)
+	}
+
+	send = time.Since(start).Seconds()
+
+	select {
+	case <-ctx.Done():
+		lm.logger.Error(ctx, "context canceled before message could be received", slog.Error(ctx.Err()), slog.F("msg", msg))
+		return send, -1, ctx.Err()
+	case val := <-res:
+		return send, val, nil
+	}
+}
+
+func (lm *LatencyMeasurer) latencyChannelName() string {
+	return fmt.Sprintf("latency-measure:%s", lm.channel)
+}

coderd/database/pubsub/pubsub.go

@@ -12,6 +12,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/lib/pq"
 	"github.com/prometheus/client_golang/prometheus"
+	"go.uber.org/atomic"
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
@@ -205,6 +206,9 @@ type PGPubsub struct {
 	receivedBytesTotal  prometheus.Counter
 	disconnectionsTotal prometheus.Counter
 	connected           prometheus.Gauge
+
+	latencyMeasurer   *LatencyMeasurer
+	latencyErrCounter atomic.Float64
 }
 
 // BufferSize is the maximum number of unhandled messages we will buffer
@@ -478,6 +482,25 @@ var (
 	)
 )
 
+// additional metrics collected out-of-band
+var (
+	pubsubSendLatencyDesc = prometheus.NewDesc(
+		"coder_pubsub_send_latency_seconds",
+		"The time taken to send a message into a pubsub event channel",
+		nil, nil,
+	)
+	pubsubRecvLatencyDesc = prometheus.NewDesc(
+		"coder_pubsub_receive_latency_seconds",
+		"The time taken to receive a message from a pubsub event channel",
+		nil, nil,
+	)
+	pubsubLatencyMeasureErrDesc = prometheus.NewDesc(
+		"coder_pubsub_latency_measure_errs_total",
+		"The number of pubsub latency measurement failures",
+		nil, nil,
+	)
+)
+
 // We'll track messages as size "normal" and "colossal", where the
 // latter are messages larger than 7600 bytes, or 95% of the postgres
 // notify limit. If we see a lot of colossal packets that's an indication that
@@ -504,6 +527,11 @@ func (p *PGPubsub) Describe(descs chan<- *prometheus.Desc) {
 	// implicit metrics
 	descs <- currentSubscribersDesc
 	descs <- currentEventsDesc
+
+	// additional metrics
+	descs <- pubsubSendLatencyDesc
+	descs <- pubsubRecvLatencyDesc
+	descs <- pubsubLatencyMeasureErrDesc
 }
 
 // Collect implements, along with Describe, the prometheus.Collector interface
@@ -528,6 +556,19 @@ func (p *PGPubsub) Collect(metrics chan<- prometheus.Metric) {
 	p.qMu.Unlock()
 	metrics <- prometheus.MustNewConstMetric(currentSubscribersDesc, prometheus.GaugeValue, float64(subs))
 	metrics <- prometheus.MustNewConstMetric(currentEventsDesc, prometheus.GaugeValue, float64(events))
+
+	// additional metrics
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
+	defer cancel()
+
+	send, recv, err := p.latencyMeasurer.Measure(ctx, p)
+	if err != nil {
+		p.logger.Warn(context.Background(), "failed to measure latency", slog.Error(err))
+		metrics <- prometheus.MustNewConstMetric(pubsubLatencyMeasureErrDesc, prometheus.CounterValue, p.latencyErrCounter.Add(1))
+		return
+	}
+	metrics <- prometheus.MustNewConstMetric(pubsubSendLatencyDesc, prometheus.GaugeValue, send)
+	metrics <- prometheus.MustNewConstMetric(pubsubRecvLatencyDesc, prometheus.GaugeValue, recv)
 }
 
 // New creates a new Pubsub implementation using a PostgreSQL connection.
@@ -544,10 +585,11 @@ func New(startCtx context.Context, logger slog.Logger, database *sql.DB, connect
 // newWithoutListener creates a new PGPubsub without creating the pqListener.
 func newWithoutListener(logger slog.Logger, database *sql.DB) *PGPubsub {
 	return &PGPubsub{
-		logger:     logger,
-		listenDone: make(chan struct{}),
-		db:         database,
-		queues:     make(map[string]map[uuid.UUID]*msgQueue),
+		logger:          logger,
+		listenDone:      make(chan struct{}),
+		db:              database,
+		queues:          make(map[string]map[uuid.UUID]*msgQueue),
+		latencyMeasurer: NewLatencyMeasurer(logger.Named("latency-measurer")),
 
 		publishesTotal: prometheus.NewCounterVec(prometheus.CounterOpts{
 			Namespace: "coder",

coderd/database/pubsub/pubsub_linux_test.go

@@ -3,11 +3,13 @@
 package pubsub_test
 
 import (
+	"bytes"
 	"context"
 	"database/sql"
 	"fmt"
 	"math/rand"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 
@@ -15,6 +17,8 @@ import (
 	"github.com/stretchr/testify/require"
 	"golang.org/x/xerrors"
 
+	"cdr.dev/slog/sloggers/sloghuman"
+
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
@@ -294,3 +298,136 @@ func TestPubsub_Disconnect(t *testing.T) {
 	}
 	require.True(t, gotDroppedErr)
 }
+
+func TestMeasureLatency(t *testing.T) {
+	t.Parallel()
+
+	newPubsub := func() (pubsub.Pubsub, func()) {
+		ctx, cancel := context.WithCancel(context.Background())
+		logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+		connectionURL, closePg, err := dbtestutil.Open()
+		require.NoError(t, err)
+		db, err := sql.Open("postgres", connectionURL)
+		require.NoError(t, err)
+		ps, err := pubsub.New(ctx, logger, db, connectionURL)
+		require.NoError(t, err)
+
+		return ps, func() {
+			_ = ps.Close()
+			_ = db.Close()
+			closePg()
+			cancel()
+		}
+	}
+
+	t.Run("MeasureLatency", func(t *testing.T) {
+		t.Parallel()
+
+		logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+		ps, done := newPubsub()
+		defer done()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+		defer cancel()
+
+		send, recv, err := pubsub.NewLatencyMeasurer(logger).Measure(ctx, ps)
+		require.NoError(t, err)
+		require.Greater(t, send, 0.0)
+		require.Greater(t, recv, 0.0)
+	})
+
+	t.Run("MeasureLatencyRecvTimeout", func(t *testing.T) {
+		t.Parallel()
+
+		logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+		ps, done := newPubsub()
+		defer done()
+
+		ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
+		defer cancel()
+
+		send, recv, err := pubsub.NewLatencyMeasurer(logger).Measure(ctx, ps)
+		require.ErrorContains(t, err, context.DeadlineExceeded.Error())
+		require.Greater(t, send, 0.0)
+		require.EqualValues(t, recv, -1)
+	})
+
+	t.Run("MeasureLatencyNotifyRace", func(t *testing.T) {
+		t.Parallel()
+
+		var buf bytes.Buffer
+		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
+		logger = logger.AppendSinks(sloghuman.Sink(&buf))
+
+		lm := pubsub.NewLatencyMeasurer(logger)
+		ps, done := newPubsub()
+		defer done()
+
+		slow := newDelayedPublisher(ps, time.Second)
+		fast := newDelayedPublisher(ps, time.Nanosecond)
+		hold := make(chan struct{}, 1)
+
+		// Start two goroutines in which two subscribers are registered but the messages are received out-of-order because
+		// the first Pubsub will publish its message slowly and the second will publish it quickly. Both will ultimately
+		// receive their desired messages, but the slow publisher will receive an unexpected message first.
+		var wg sync.WaitGroup
+		wg.Add(2)
+		go func() {
+			ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+			defer cancel()
+			defer wg.Done()
+
+			hold <- struct{}{}
+			send, recv, err := lm.Measure(ctx, slow)
+			assert.NoError(t, err)
+			assert.Greater(t, send, 0.0)
+			assert.Greater(t, recv, 0.0)
+
+			// The slow subscriber will complete first and receive the fast publisher's message first.
+			logger.Sync()
+			assert.Contains(t, buf.String(), "received unexpected message")
+		}()
+
+		go func() {
+			ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+			defer cancel()
+			defer wg.Done()
+
+			// Force fast publisher to start after the slow one to avoid flakiness.
+			<-hold
+			time.Sleep(time.Millisecond * 50)
+			send, recv, err := lm.Measure(ctx, fast)
+			assert.NoError(t, err)
+			assert.Greater(t, send, 0.0)
+			assert.Greater(t, recv, 0.0)
+		}()
+
+		wg.Wait()
+	})
+}
+
+type delayedPublisher struct {
+	pubsub.Pubsub
+	delay time.Duration
+}
+
+func newDelayedPublisher(ps pubsub.Pubsub, delay time.Duration) *delayedPublisher {
+	return &delayedPublisher{Pubsub: ps, delay: delay}
+}
+
+func (s *delayedPublisher) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
+	return s.Pubsub.Subscribe(event, listener)
+}
+
+func (s *delayedPublisher) SubscribeWithErr(event string, listener pubsub.ListenerWithErr) (cancel func(), err error) {
+	return s.Pubsub.SubscribeWithErr(event, listener)
+}
+
+func (s *delayedPublisher) Publish(event string, message []byte) error {
+	time.Sleep(s.delay)
+	return s.Pubsub.Publish(event, message)
+}
+
+func (s *delayedPublisher) Close() error {
+	return s.Pubsub.Close()
+}

coderd/database/pubsub/pubsub_test.go

@@ -39,7 +39,11 @@ func TestPGPubsub_Metrics(t *testing.T) {
 	err = registry.Register(uut)
 	require.NoError(t, err)
 
+	// each Gather measures pubsub latency by publishing a message & subscribing to it
+	var gatherCount float64
+
 	metrics, err := registry.Gather()
+	gatherCount++
 	require.NoError(t, err)
 	require.True(t, testutil.PromGaugeHasValue(t, metrics, 0, "coder_pubsub_current_events"))
 	require.True(t, testutil.PromGaugeHasValue(t, metrics, 0, "coder_pubsub_current_subscribers"))
@@ -59,19 +63,25 @@ func TestPGPubsub_Metrics(t *testing.T) {
 	_ = testutil.RequireRecvCtx(ctx, t, messageChannel)
 
 	require.Eventually(t, func() bool {
+		latencyBytes := gatherCount * pubsub.LatencyMessageLength
 		metrics, err = registry.Gather()
+		gatherCount++
 		assert.NoError(t, err)
 		return testutil.PromGaugeHasValue(t, metrics, 1, "coder_pubsub_current_events") &&
 			testutil.PromGaugeHasValue(t, metrics, 1, "coder_pubsub_current_subscribers") &&
 			testutil.PromGaugeHasValue(t, metrics, 1, "coder_pubsub_connected") &&
-			testutil.PromCounterHasValue(t, metrics, 1, "coder_pubsub_publishes_total", "true") &&
-			testutil.PromCounterHasValue(t, metrics, 1, "coder_pubsub_subscribes_total", "true") &&
-			testutil.PromCounterHasValue(t, metrics, 1, "coder_pubsub_messages_total", "normal") &&
-			testutil.PromCounterHasValue(t, metrics, 7, "coder_pubsub_received_bytes_total") &&
-			testutil.PromCounterHasValue(t, metrics, 7, "coder_pubsub_published_bytes_total")
+			testutil.PromCounterHasValue(t, metrics, gatherCount, "coder_pubsub_publishes_total", "true") &&
+			testutil.PromCounterHasValue(t, metrics, gatherCount, "coder_pubsub_subscribes_total", "true") &&
+			testutil.PromCounterHasValue(t, metrics, gatherCount, "coder_pubsub_messages_total", "normal") &&
+			testutil.PromCounterHasValue(t, metrics, float64(len(data))+latencyBytes, "coder_pubsub_received_bytes_total") &&
+			testutil.PromCounterHasValue(t, metrics, float64(len(data))+latencyBytes, "coder_pubsub_published_bytes_total") &&
+			testutil.PromGaugeAssertion(t, metrics, func(in float64) bool { return in > 0 }, "coder_pubsub_send_latency_seconds") &&
+			testutil.PromGaugeAssertion(t, metrics, func(in float64) bool { return in > 0 }, "coder_pubsub_receive_latency_seconds") &&
+			!testutil.PromCounterGathered(t, metrics, "coder_pubsub_latency_measure_errs_total")
 	}, testutil.WaitShort, testutil.IntervalFast)
 
-	colossalData := make([]byte, 7600)
+	colossalSize := 7600
+	colossalData := make([]byte, colossalSize)
 	for i := range colossalData {
 		colossalData[i] = 'q'
 	}
@@ -89,16 +99,21 @@ func TestPGPubsub_Metrics(t *testing.T) {
 	_ = testutil.RequireRecvCtx(ctx, t, messageChannel)
 
 	require.Eventually(t, func() bool {
+		latencyBytes := gatherCount * pubsub.LatencyMessageLength
 		metrics, err = registry.Gather()
+		gatherCount++
 		assert.NoError(t, err)
 		return testutil.PromGaugeHasValue(t, metrics, 1, "coder_pubsub_current_events") &&
 			testutil.PromGaugeHasValue(t, metrics, 2, "coder_pubsub_current_subscribers") &&
 			testutil.PromGaugeHasValue(t, metrics, 1, "coder_pubsub_connected") &&
-			testutil.PromCounterHasValue(t, metrics, 2, "coder_pubsub_publishes_total", "true") &&
-			testutil.PromCounterHasValue(t, metrics, 2, "coder_pubsub_subscribes_total", "true") &&
-			testutil.PromCounterHasValue(t, metrics, 1, "coder_pubsub_messages_total", "normal") &&
+			testutil.PromCounterHasValue(t, metrics, 1+gatherCount, "coder_pubsub_publishes_total", "true") &&
+			testutil.PromCounterHasValue(t, metrics, 1+gatherCount, "coder_pubsub_subscribes_total", "true") &&
+			testutil.PromCounterHasValue(t, metrics, gatherCount, "coder_pubsub_messages_total", "normal") &&
 			testutil.PromCounterHasValue(t, metrics, 1, "coder_pubsub_messages_total", "colossal") &&
-			testutil.PromCounterHasValue(t, metrics, 7607, "coder_pubsub_received_bytes_total") &&
-			testutil.PromCounterHasValue(t, metrics, 7607, "coder_pubsub_published_bytes_total")
+			testutil.PromCounterHasValue(t, metrics, float64(colossalSize+len(data))+latencyBytes, "coder_pubsub_received_bytes_total") &&
+			testutil.PromCounterHasValue(t, metrics, float64(colossalSize+len(data))+latencyBytes, "coder_pubsub_published_bytes_total") &&
+			testutil.PromGaugeAssertion(t, metrics, func(in float64) bool { return in > 0 }, "coder_pubsub_send_latency_seconds") &&
+			testutil.PromGaugeAssertion(t, metrics, func(in float64) bool { return in > 0 }, "coder_pubsub_receive_latency_seconds") &&
+			!testutil.PromCounterGathered(t, metrics, "coder_pubsub_latency_measure_errs_total")
 	}, testutil.WaitShort, testutil.IntervalFast)
 }

scripts/linux-pkg/coder-openrc

@@ -0,0 +1,38 @@
+#!/sbin/openrc-run
+name=coder
+description="Coder - Self-hosted developer workspaces on your infra"
+document="https://coder.com/docs/coder-oss"
+
+depend() {
+    need net
+    after net-online
+    use dns logger
+}
+
+checkpath --directory --owner coder:coder --mode 0700 /var/cache/coder
+
+start_pre() {
+    if [ ! -f /etc/coder.d/coder.env ]; then
+        eerror "/etc/coder.d/coder.env file does not exist"
+        return 1
+    fi
+    # Read and export environment variables ignoring comment lines and blank lines
+    while IFS= read -r line; do
+        # Skip blank or comment lines
+        if [ -z "$line" ] || [[ "$line" =~ ^# ]]; then
+            continue
+        fi
+        export "$line"
+    done < /etc/coder.d/coder.env
+}
+
+command="/usr/bin/coder"
+command_args="server"
+command_user="coder:coder"
+command_background="yes"
+pidfile="/run/coder.pid"
+
+restart="always"
+restart_delay="5"
+
+stop_timeout="90"

scripts/linux-pkg/coder-workspace-proxy-openrc

@@ -0,0 +1,39 @@
+#!/sbin/openrc-run
+name=coder-workspace-proxy
+description="Coder - external workspace proxy server"
+document="https://coder.com/docs/coder-oss"
+
+depend() {
+    need net
+    after net-online
+    use dns logger
+}
+
+checkpath --directory --owner coder:coder --mode 0700 /var/cache/coder
+
+start_pre() {
+	if [ ! -f /etc/coder.d/coder-workspace-proxy.env ]; then
+		eerror "/etc/coder.d/coder-workspace-proxy.env file does not exist"
+		return 1
+	fi
+
+    # Read and export environment variables ignoring comment lines and blank lines
+    while IFS= read -r line; do
+        # Skip blank or comment lines
+        if [ -z "$line" ] || [[ "$line" =~ ^# ]]; then
+            continue
+        fi
+        export "$line"
+    done < /etc/coder.d/coder-workspace-proxy.env
+}
+
+command="/usr/bin/coder"
+command_args="workspace-proxy server"
+command_user="coder:coder"
+command_background="yes"
+pidfile="/run/coder-workspace-proxy.pid"
+
+restart="always"
+restart_delay="5"
+
+stop_timeout="90"

scripts/linux-pkg/nfpm-alpine.yaml

@@ -0,0 +1,29 @@
+name: coder
+platform: linux
+arch: "${GOARCH}"
+version: "${CODER_VERSION}"
+version_schema: semver
+release: 1
+
+vendor: Coder
+homepage: https://coder.com
+maintainer: Coder <support@coder.com>
+description: |
+  Provision development environments with infrastructure with code
+license: AGPL-3.0
+suggests:
+  - postgresql
+
+scripts:
+  preinstall: preinstall.sh
+
+contents:
+  - src: coder
+    dst: /usr/bin/coder
+  - src: coder.env
+    dst: /etc/coder.d/coder.env
+    type: "config|noreplace"
+  - src: coder-workspace-proxy-openrc
+    dst: /etc/init.d/coder-workspace-proxy
+  - src: coder-openrc
+    dst: /etc/init.d/coder

scripts/package.sh

@@ -89,9 +89,16 @@ ln "$(realpath scripts/linux-pkg/coder.service)" "$temp_dir/"
 ln "$(realpath scripts/linux-pkg/nfpm.yaml)" "$temp_dir/"
 ln "$(realpath scripts/linux-pkg/preinstall.sh)" "$temp_dir/"
 
+nfpm_config_file="nfpm.yaml"
+
+# Use nfpm-alpine.yaml when building for Alpine (OpenRC).
+if [[ "$format" == "apk" ]]; then
+	nfpm_config_file="nfpm-alpine.yaml"
+fi
+
 pushd "$temp_dir"
 GOARCH="$arch" CODER_VERSION="$version" nfpm package \
-	-f nfpm.yaml \
+	-f "$nfpm_config_file" \
 	-p "$format" \
 	-t "$output_path" \
 	1>&2

testutil/prometheus.go

@@ -7,29 +7,58 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func PromGaugeHasValue(t testing.TB, metrics []*dto.MetricFamily, value float64, name string, label ...string) bool {
+type kind string
+
+const (
+	counterKind kind = "counter"
+	gaugeKind   kind = "gauge"
+)
+
+func PromGaugeHasValue(t testing.TB, metrics []*dto.MetricFamily, value float64, name string, labels ...string) bool {
 	t.Helper()
-	for _, family := range metrics {
-		if family.GetName() != name {
-			continue
-		}
-		ms := family.GetMetric()
-	metricsLoop:
-		for _, m := range ms {
-			require.Equal(t, len(label), len(m.GetLabel()))
-			for i, lv := range label {
-				if lv != m.GetLabel()[i].GetValue() {
-					continue metricsLoop
-				}
-			}
-			return value == m.GetGauge().GetValue()
-		}
-	}
-	return false
+	return value == getValue(t, metrics, gaugeKind, name, labels...)
 }
 
-func PromCounterHasValue(t testing.TB, metrics []*dto.MetricFamily, value float64, name string, label ...string) bool {
+func PromCounterHasValue(t testing.TB, metrics []*dto.MetricFamily, value float64, name string, labels ...string) bool {
 	t.Helper()
+	return value == getValue(t, metrics, counterKind, name, labels...)
+}
+
+func PromGaugeAssertion(t testing.TB, metrics []*dto.MetricFamily, assert func(in float64) bool, name string, labels ...string) bool {
+	t.Helper()
+	return assert(getValue(t, metrics, gaugeKind, name, labels...))
+}
+
+func PromCounterAssertion(t testing.TB, metrics []*dto.MetricFamily, assert func(in float64) bool, name string, labels ...string) bool {
+	t.Helper()
+	return assert(getValue(t, metrics, counterKind, name, labels...))
+}
+
+func PromCounterGathered(t testing.TB, metrics []*dto.MetricFamily, name string, labels ...string) bool {
+	return getMetric(t, metrics, name, labels...) != nil
+}
+
+func PromGaugeGathered(t testing.TB, metrics []*dto.MetricFamily, name string, labels ...string) bool {
+	return getMetric(t, metrics, name, labels...) != nil
+}
+
+func getValue(t testing.TB, metrics []*dto.MetricFamily, kind kind, name string, labels ...string) float64 {
+	m := getMetric(t, metrics, name, labels...)
+	if m == nil {
+		return -1
+	}
+
+	switch kind {
+	case counterKind:
+		return m.GetCounter().GetValue()
+	case gaugeKind:
+		return m.GetGauge().GetValue()
+	default:
+		return -1
+	}
+}
+
+func getMetric(t testing.TB, metrics []*dto.MetricFamily, name string, labels ...string) *dto.Metric {
 	for _, family := range metrics {
 		if family.GetName() != name {
 			continue
@@ -37,14 +66,16 @@ func PromCounterHasValue(t testing.TB, metrics []*dto.MetricFamily, value float6
 		ms := family.GetMetric()
 	metricsLoop:
 		for _, m := range ms {
-			require.Equal(t, len(label), len(m.GetLabel()))
-			for i, lv := range label {
+			require.Equal(t, len(labels), len(m.GetLabel()))
+			for i, lv := range labels {
 				if lv != m.GetLabel()[i].GetValue() {
 					continue metricsLoop
 				}
 			}
-			return value == m.GetCounter().GetValue()
+
+			return m
 		}
 	}
-	return false
+
+	return nil
 }