relay: the pprof upstream kind

cmd/perfloop-relay/main_test.go

@@ -25,7 +25,7 @@ func TestPrintRoutesPrintsTheTableOrValidatesError(t *testing.T) {
 	}
 	bad := filepath.Join(t.TempDir(), "relay.yaml")
 	for body, want := range map[string]string{
-		"api: https://app.perfloop.ai\ntoken: x\nupstreams:\n  - name: vm\n    kind: pprof\n    url: http://vm.internal\n":                           `kind "pprof" is not supported`,
+		"api: https://app.perfloop.ai\ntoken: x\nupstreams:\n  - name: vm\n    kind: pyroscope\n    url: http://vm.internal\n":                       `kind "pyroscope" is not supported`,
 		"api: https://app.perfloop.ai\ntoken: x\nlog_level: nonsense\nupstreams:\n  - name: vm\n    kind: prometheus\n    url: http://vm.internal\n": `log_level:`,
 	} {
 		if err := os.WriteFile(bad, []byte(body), 0o600); err != nil {

config.go

@@ -41,12 +41,28 @@ type Upstream struct {
 	// lowercase DNS label.
 	Name string `yaml:"name"`
 	// Kind selects the read routes the relay permits; routes.go lists the
-	// kinds (today prometheus, victoriametrics, and loki). It selects the read
-	// routes the relay permits.
+	// kinds (today prometheus, victoriametrics, loki, and pprof).
 	Kind string `yaml:"kind"`
 	// URL is the provider base: scheme, host, port, and an optional path
-	// prefix. The relay appends the proxy's request path to it.
+	// prefix. The relay appends the proxy's request path to it. A pprof
+	// upstream has Targets instead.
 	URL string `yaml:"url"`
+	// Targets are the processes of a pprof upstream: target name to the
+	// http(s) base of that process's net/http/pprof handler. Perfloop names
+	// one target per read as the first path segment; a name not listed here
+	// is refused. Each name is one lowercase DNS label.
+	Targets map[string]string `yaml:"targets"`
+	// MaxSeconds caps the `seconds` a pprof read may ask for, so a CPU
+	// profile cannot run longer than the customer allows. Zero means
+	// DefaultMaxSeconds. The bound is link.MaxProfileSeconds.
+	MaxSeconds int `yaml:"max_seconds"`
+	// MaxConcurrent caps the pprof reads in flight on this upstream across
+	// all of its targets, per relay process; a further read is refused at
+	// once with status 429. Perfloop spreads reads over every relay a tenant
+	// runs with the same token, so N relay processes admit N times this cap.
+	// The customer controls that count. Zero means DefaultMaxConcurrent. The
+	// bound is MaxConcurrentBound.
+	MaxConcurrent int `yaml:"max_concurrent"`
 	// Headers are added to every upstream request, for example an
 	// Authorization header holding a read token. They replace any header of
 	// the same name from the proxy.
@@ -56,6 +72,20 @@ type Upstream struct {
 	CA string `yaml:"ca"`
 }
 
+const (
+	// DefaultMaxSeconds is the `seconds` cap of a pprof upstream that sets
+	// none: net/http/pprof's own default CPU profile length. The largest cap
+	// a config may set is link.MaxProfileSeconds.
+	DefaultMaxSeconds = 30
+	// DefaultMaxConcurrent is the in-flight cap of a pprof upstream that
+	// sets none: one profile at a time per upstream.
+	DefaultMaxConcurrent = 1
+	// MaxConcurrentBound is the largest max_concurrent a config may set. It
+	// is a config bound on how many processes profile at once, not a tunnel
+	// property: the relay keeps two tunnels of maxStreams each.
+	MaxConcurrentBound = 8
+)
+
 // Load reads, expands, and strictly parses a relay configuration. New
 // validates it.
 func Load(path string) (Config, error) {
@@ -102,9 +132,34 @@ func (c Config) Validate() error {
 		if _, ok := base(up.Kind); !ok {
 			return fmt.Errorf("upstream %s: kind %q is not supported", up.Name, up.Kind)
 		}
-		u, err := url.Parse(up.URL)
-		if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" || strings.HasSuffix(u.Path, "/") {
-			return fmt.Errorf("upstream %s: url must be an http(s) base without query, fragment, user info, or trailing slash", up.Name)
+		if up.Kind == "pprof" {
+			if up.URL != "" {
+				return fmt.Errorf("upstream %s: a pprof upstream names targets, not a url", up.Name)
+			}
+			if len(up.Targets) == 0 {
+				return fmt.Errorf("upstream %s: at least one target is required", up.Name)
+			}
+			for target, raw := range up.Targets {
+				if !link.ValidUpstreamName(target) {
+					return fmt.Errorf("upstream %s: target name %q must be one lowercase DNS label", up.Name, target)
+				}
+				if !validBase(raw) {
+					return fmt.Errorf("upstream %s: target %s: url must be an http(s) base without query, fragment, user info, or trailing slash", up.Name, target)
+				}
+			}
+			if up.MaxSeconds < 0 || up.MaxSeconds > link.MaxProfileSeconds {
+				return fmt.Errorf("upstream %s: max_seconds must be between 1 and %d", up.Name, link.MaxProfileSeconds)
+			}
+			if up.MaxConcurrent < 0 || up.MaxConcurrent > MaxConcurrentBound {
+				return fmt.Errorf("upstream %s: max_concurrent must be between 1 and %d", up.Name, MaxConcurrentBound)
+			}
+		} else {
+			if up.Targets != nil || up.MaxSeconds != 0 || up.MaxConcurrent != 0 {
+				return fmt.Errorf("upstream %s: targets, max_seconds, and max_concurrent apply only to kind pprof", up.Name)
+			}
+			if !validBase(up.URL) {
+				return fmt.Errorf("upstream %s: url must be an http(s) base without query, fragment, user info, or trailing slash", up.Name)
+			}
 		}
 		// Header names and values must be ones the HTTP client will send, and
 		// one wire header must have one value: `Authorization` and
@@ -142,6 +197,13 @@ func (c Config) Level() slog.Level {
 	return level
 }
 
+// validBase reports whether raw is an http(s) base the relay may append a
+// route to: scheme, host, port, and an optional path prefix, nothing else.
+func validBase(raw string) bool {
+	u, err := url.Parse(raw)
+	return err == nil && (u.Scheme == "http" || u.Scheme == "https") && u.Host != "" && u.User == nil && u.RawQuery == "" && u.Fragment == "" && !strings.HasSuffix(u.Path, "/")
+}
+
 // apiURL is the API origin: an https URL with a host and nothing else. The
 // token travels only inside TLS.
 func (c Config) apiURL() (*url.URL, error) {

config_test.go

@@ -5,6 +5,8 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+
+	"github.com/perfloop/relay/link"
 )
 
 func TestLoadExpandsSecretsAndValidates(t *testing.T) {
@@ -29,6 +31,13 @@ upstreams:
     kind: loki
     url: https://loki-gateway.monitoring.svc
     ca: /etc/perfloop-relay/loki-ca.pem
+  - name: go
+    kind: pprof
+    targets:
+      api-1: http://10.0.0.1:6060
+      api-2: http://10.0.0.2:6060/internal
+    max_seconds: 10
+    max_concurrent: 2
 `)
 	cfg, err := Load(path)
 	if err != nil {
@@ -37,6 +46,12 @@ upstreams:
 	if cfg.Token != "tenant-seven" || cfg.Upstreams[0].Headers["Authorization"] != "Bearer vm-read" || cfg.Upstreams[1].CA == "" {
 		t.Fatalf("loaded %+v", cfg)
 	}
+	if go1 := cfg.Upstreams[2]; len(go1.Targets) != 2 || go1.Targets["api-2"] != "http://10.0.0.2:6060/internal" || go1.MaxSeconds != 10 || go1.MaxConcurrent != 2 {
+		t.Fatalf("loaded pprof upstream %+v", go1)
+	}
+	if err := cfg.Validate(); err != nil {
+		t.Fatal(err)
+	}
 	write("api: https://app.perfloop.ai\ntoken: ${RELAY_TEST_MISSING}\nupstreams: [{name: vm, kind: prometheus, url: http://vm:8428}]\n")
 	if _, err := Load(path); err == nil || !strings.Contains(err.Error(), "RELAY_TEST_MISSING") {
 		t.Fatalf("unset variable: %v", err)
@@ -77,6 +92,32 @@ upstreams:
 		"one header spelled twice": func(c *Config) {
 			c.Upstreams[0].Headers = map[string]string{"Authorization": "Bearer x", "authorization": "Bearer y"}
 		},
+		"targets on a metrics upstream":     func(c *Config) { c.Upstreams[0].Targets = map[string]string{"a": "http://a:6060"} },
+		"max_seconds on a metrics upstream": func(c *Config) { c.Upstreams[0].MaxSeconds = 10 },
+		"pprof with url":                    func(c *Config) { c.Upstreams[0] = pprof(func(u *Upstream) { u.URL = "http://a:6060" }) },
+		"pprof without targets":             func(c *Config) { c.Upstreams[0] = pprof(func(u *Upstream) { u.Targets = nil }) },
+		"pprof target name with a dot": func(c *Config) {
+			c.Upstreams[0] = pprof(func(u *Upstream) { u.Targets = map[string]string{"api.1": "http://a:6060"} })
+		},
+		"pprof target name with a slash": func(c *Config) {
+			c.Upstreams[0] = pprof(func(u *Upstream) { u.Targets = map[string]string{"api/1": "http://a:6060"} })
+		},
+		"pprof target url with query": func(c *Config) {
+			c.Upstreams[0] = pprof(func(u *Upstream) { u.Targets = map[string]string{"a": "http://a:6060/?x=1"} })
+		},
+		"pprof target url with user": func(c *Config) {
+			c.Upstreams[0] = pprof(func(u *Upstream) { u.Targets = map[string]string{"a": "http://u:p@a:6060"} })
+		},
+		"pprof target url with trailing slash": func(c *Config) {
+			c.Upstreams[0] = pprof(func(u *Upstream) { u.Targets = map[string]string{"a": "http://a:6060/"} })
+		},
+		"pprof max_seconds over the bound": func(c *Config) {
+			c.Upstreams[0] = pprof(func(u *Upstream) { u.MaxSeconds = link.MaxProfileSeconds + 1 })
+		},
+		"pprof max_seconds negative": func(c *Config) { c.Upstreams[0] = pprof(func(u *Upstream) { u.MaxSeconds = -1 }) },
+		"pprof max_concurrent over the bound": func(c *Config) {
+			c.Upstreams[0] = pprof(func(u *Upstream) { u.MaxConcurrent = MaxConcurrentBound + 1 })
+		},
 	} {
 		cfg := base
 		cfg.Upstreams = append([]Upstream(nil), base.Upstreams...)
@@ -86,3 +127,17 @@ upstreams:
 		}
 	}
 }
+
+// pprof is a valid pprof upstream with one change applied.
+func pprof(change func(*Upstream)) Upstream {
+	up := Upstream{Name: "go", Kind: "pprof", Targets: map[string]string{"api-1": "http://10.0.0.1:6060"}}
+	change(&up)
+	return up
+}
+
+func TestValidPprofUpstream(t *testing.T) {
+	cfg := Config{API: "https://app.perfloop.ai", Token: "t", Upstreams: []Upstream{pprof(func(*Upstream) {})}}
+	if err := cfg.Validate(); err != nil {
+		t.Fatal(err)
+	}
+}

link/link.go

@@ -37,6 +37,25 @@ const (
 	// it, its tunnel token covers it, and the relay logs it and drops it
 	// before the upstream request. It is absent on a read no tool call made.
 	RequestIDHeader = "Perfloop-Request-Id"
+	// ErrorHeader marks a response that is not the provider's: Perfloop made
+	// it (Unavailable) or the relay refused the read (Refused). A reader can
+	// tell either from provider bytes; the API and the proxy act only on
+	// Unavailable, and a relay refusal reaches the adapter with its status.
+	ErrorHeader = "Perfloop-Relay-Error"
+	// Unavailable is the value Perfloop sets when no relay could serve the
+	// read: no tunnel, or a token the replica does not recognize.
+	Unavailable = "1"
+	// Refused is the value the relay sets on its own refusals: an unknown
+	// upstream or target, a route or query outside the allowlist, or the
+	// in-flight cap. The status says which.
+	Refused = "relay"
+	// MaxProfileSeconds is the longest `seconds` a pprof read may ask for on
+	// either end: the relay refuses a larger `max_seconds` in its config, and
+	// Perfloop's pprof adapter refuses a longer capture before the tunnel. A
+	// profile read crosses two one-minute deadlines, the controller's remote
+	// client and the proxy's read; 45 leaves them 15 seconds for the tunnel
+	// and the bytes.
+	MaxProfileSeconds = 45
 )
 
 // hopHeaders never cross a tunnel in either direction (RFC 9110 § 7.6.1),

relay.go

@@ -7,6 +7,7 @@
 package relay
 
 import (
+	"cmp"
 	"context"
 	"crypto/tls"
 	"crypto/x509"
@@ -73,8 +74,14 @@ type Relay struct {
 
 type upstream struct {
 	Upstream
-	base   *url.URL
-	client *http.Client
+	// base is the provider URL of a metrics or logs upstream.
+	base *url.URL
+	// targets are the process bases of a pprof upstream, by target name.
+	targets map[string]*url.URL
+	// slots is the pprof in-flight cap: one token per permitted read.
+	slots      chan struct{}
+	maxSeconds int
+	client     *http.Client
 }
 
 // New validates the configuration and builds a relay.
@@ -126,9 +133,16 @@ func New(cfg Config, logger *slog.Logger) (*Relay, error) {
 			IdleConnTimeout:       90 * time.Second,
 			ResponseHeaderTimeout: readTimeout,
 		}
+		targets := make(map[string]*url.URL, len(up.Targets))
+		for name, raw := range up.Targets {
+			targets[name], _ = url.Parse(raw)
+		}
 		r.upstreams[up.Name] = &upstream{
-			Upstream: up,
-			base:     base,
+			Upstream:   up,
+			base:       base,
+			targets:    targets,
+			slots:      make(chan struct{}, cmp.Or(up.MaxConcurrent, DefaultMaxConcurrent)),
+			maxSeconds: cmp.Or(up.MaxSeconds, DefaultMaxSeconds),
 			client: &http.Client{Transport: upstreamTransport, Timeout: readTimeout,
 				CheckRedirect: func(*http.Request, []*http.Request) error { return errors.New("upstream redirects are not followed") }},
 		}
@@ -224,7 +238,9 @@ func (r *Relay) serve(ctx context.Context) error {
 
 // ServeHTTP answers one proxy read from the tunnel. The Host names the
 // upstream; the path is forwarded under the upstream URL when its route is a
-// documented read. One audit line records every decision with the read id
+// documented read. For a pprof upstream the first path segment names the
+// target process and the query is checked against the configured caps. One
+// audit line records every decision with the read id
 // Perfloop gave it, `<session id>/<tool call ref>`, which names the
 // transcript turn that asked; a read no tool call made has none.
 func (r *Relay) ServeHTTP(w http.ResponseWriter, req *http.Request) {
@@ -250,10 +266,35 @@ func (r *Relay) forward(w http.ResponseWriter, req *http.Request) (int, string,
 	if p := req.URL.Path; p == "" || p[0] != '/' || p != path.Clean(p) {
 		return refuse(w, http.StatusBadRequest, "path is not clean"), "path", nil
 	}
-	name, ok := route(up.Kind, req.URL.Path)
+	base, rest, query := up.base, req.URL.Path, req.URL.RawQuery
+	if up.Kind == "pprof" {
+		// The target and the query are checked here, before any request, in
+		// the same place as the route: the target map and the caps are the
+		// customer's control over which process runs a profile and for how
+		// long. The query forwarded is the one checked, not the raw bytes.
+		var name string
+		var ok bool
+		if name, rest, ok = target(req.URL.Path); !ok || up.targets[name] == nil {
+			return refuse(w, http.StatusNotFound, "unknown target"), "unknown-target", nil
+		}
+		base = up.targets[name]
+	}
+	name, ok := route(up.Kind, rest)
 	if !ok {
 		return refuse(w, http.StatusForbidden, "route is not a documented read"), "route", nil
 	}
+	if up.Kind == "pprof" {
+		var err error
+		if query, err = pprofQuery(name, req.URL.RawQuery, up.maxSeconds); err != nil {
+			return refuse(w, http.StatusBadRequest, err.Error()), "query", nil
+		}
+		select {
+		case up.slots <- struct{}{}:
+			defer func() { <-up.slots }()
+		default:
+			return refuse(w, http.StatusTooManyRequests, "profile reads in flight are at the configured cap"), "busy", nil
+		}
+	}
 	// One read is bounded end to end on both sides: the upstream client's
 	// timeout, and the same deadline on the response write, so a tunnel peer
 	// that withholds flow-control credit cannot hold the stream forever.
@@ -262,9 +303,9 @@ func (r *Relay) forward(w http.ResponseWriter, req *http.Request) (int, string,
 	if err := http.NewResponseController(w).SetWriteDeadline(time.Now().Add(readTimeout)); err != nil {
 		return refuse(w, http.StatusInternalServerError, "write deadline is not supported"), "path", err
 	}
-	target := *up.base
-	target.Path = path.Join(up.base.Path, req.URL.Path)
-	target.RawQuery = req.URL.RawQuery
+	target := *base
+	target.Path = path.Join(base.Path, rest)
+	target.RawQuery = query
 	out, err := http.NewRequestWithContext(ctx, http.MethodGet, target.String(), http.NoBody)
 	if err != nil {
 		return refuse(w, http.StatusBadRequest, "request could not be built"), "path", err
@@ -310,7 +351,11 @@ func (r *Relay) forward(w http.ResponseWriter, req *http.Request) (int, string,
 	return resp.StatusCode, name, nil
 }
 
+// refuse answers a read the relay will not forward. The error header names
+// the relay as the author, so a reader can tell this 404 or 400 from one the
+// upstream itself returned.
 func refuse(w http.ResponseWriter, status int, reason string) int {
+	w.Header().Set(link.ErrorHeader, link.Refused)
 	http.Error(w, reason, status)
 	return status
 }