From 2581e387899413e9933d28101a1a3707331f0327 Mon Sep 17 00:00:00 2001
From: Mike O'Driscoll <mikeo@tailscale.com>
Date: Tue, 19 Aug 2025 12:13:55 -0400
Subject: [PATCH] prober: update runall handler to be generic (#16895)

Update the runall handler to be more generic with an
exclude param to exclude multiple probes as the requesters
definition.

Updates tailscale/corp#27370

Signed-off-by: Mike O'Driscoll <mikeo@tailscale.com>
---
 prober/prober.go      |  5 +++-
 prober/prober_test.go | 69 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/prober/prober.go b/prober/prober.go
index b69d26821..9c494c3c9 100644
--- a/prober/prober.go
+++ b/prober/prober.go
@@ -18,6 +18,7 @@ import (
 	"maps"
 	"math/rand"
 	"net/http"
+	"slices"
 	"sync"
 	"time"
 
@@ -585,10 +586,12 @@ type RunHandlerAllResponse struct {
 }
 
 func (p *Prober) RunAllHandler(w http.ResponseWriter, r *http.Request) error {
+	excluded := r.URL.Query()["exclude"]
+
 	probes := make(map[string]*Probe)
 	p.mu.Lock()
 	for _, probe := range p.probes {
-		if !probe.IsContinuous() && probe.name != "derpmap-probe" {
+		if !probe.IsContinuous() && !slices.Contains(excluded, probe.name) {
 			probes[probe.name] = probe
 		}
 	}
diff --git a/prober/prober_test.go b/prober/prober_test.go
index 7cb841936..15db21a5e 100644
--- a/prober/prober_test.go
+++ b/prober/prober_test.go
@@ -11,6 +11,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -722,7 +723,7 @@ func TestRunAllHandler(t *testing.T) {
 
 			mux.Handle("/prober/runall/", tsweb.StdHandler(tsweb.ReturnHandlerFunc(p.RunAllHandler), tsweb.HandlerOptions{}))
 
-			req, err := http.NewRequest("GET", server.URL+"/prober/runall/", nil)
+			req, err := http.NewRequest("GET", server.URL+"/prober/runall", nil)
 			if err != nil {
 				t.Fatalf("failed to create request: %v", err)
 			}
@@ -757,6 +758,72 @@ func TestRunAllHandler(t *testing.T) {
 
 }
 
+func TestExcludeInRunAll(t *testing.T) {
+	clk := newFakeTime()
+	p := newForTest(clk.Now, clk.NewTicker).WithOnce(true)
+
+	wantJSONResponse := RunHandlerAllResponse{
+		Results: map[string]RunHandlerResponse{
+			"includedProbe": {
+				ProbeInfo: ProbeInfo{
+					Name:          "includedProbe",
+					Interval:      probeInterval,
+					Status:        ProbeStatusSucceeded,
+					RecentResults: []bool{true, true},
+				},
+				PreviousSuccessRatio: 1,
+			},
+		},
+	}
+
+	p.Run("includedProbe", probeInterval, nil, FuncProbe(func(context.Context) error { return nil }))
+	p.Run("excludedProbe", probeInterval, nil, FuncProbe(func(context.Context) error { return nil }))
+	p.Run("excludedOtherProbe", probeInterval, nil, FuncProbe(func(context.Context) error { return nil }))
+
+	mux := http.NewServeMux()
+	server := httptest.NewServer(mux)
+	defer server.Close()
+
+	mux.Handle("/prober/runall", tsweb.StdHandler(tsweb.ReturnHandlerFunc(p.RunAllHandler), tsweb.HandlerOptions{}))
+
+	req, err := http.NewRequest("GET", server.URL+"/prober/runall", nil)
+	if err != nil {
+		t.Fatalf("failed to create request: %v", err)
+	}
+
+	// Exclude probes with "excluded" in their name
+	req.URL.RawQuery = url.Values{
+		"exclude": []string{"excludedProbe", "excludedOtherProbe"},
+	}.Encode()
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatalf("failed to make request: %v", err)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		t.Errorf("unexpected response code: got %d, want %d", resp.StatusCode, http.StatusOK)
+	}
+
+	var gotJSON RunHandlerAllResponse
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatalf("failed to read response body: %v", err)
+	}
+
+	if err := json.Unmarshal(body, &gotJSON); err != nil {
+		t.Fatalf("failed to unmarshal JSON response: %v; body: %s", err, body)
+	}
+
+	if resp.Header.Get("Content-Type") != "application/json" {
+		t.Errorf("unexpected content type: got %q, want application/json", resp.Header.Get("Content-Type"))
+	}
+
+	if diff := cmp.Diff(wantJSONResponse, gotJSON, cmpopts.IgnoreFields(ProbeInfo{}, "Start", "End", "Labels", "RecentLatencies")); diff != "" {
+		t.Errorf("unexpected JSON response (-want +got):\n%s", diff)
+	}
+}
+
 type fakeTicker struct {
 	ch       chan time.Time
 	interval time.Duration