feat: relay to phlare

README.md

@@ -23,15 +23,19 @@ If needed, the `PYROSCOPE_AUTH_TOKEN` can be supplied.
 For a complete list of variables check the section below.
 
 ## Configuration
-| env var                      | default                          | description                                                                     |
-|------------------------------|----------------------------------|---------------------------------------------------------------------------------|
-| `PYROSCOPE_REMOTE_ADDRESS`   | `https://ingest.pyroscope.cloud` | the pyroscope instance data will be relayed to                                  |
-| `PYROSCOPE_AUTH_TOKEN`       | `""`                             | authorization key (token authentication)                                        |
-| `PYROSCOPE_SELF_PROFILING`   | `false`                          | whether to profile the extension itself or not                                  |
-| `PYROSCOPE_LOG_LEVEL`        | `info`                           | `error` or `info` or `debug` or `trace`                                         |
-| `PYROSCOPE_TIMEOUT`          | `10s`                            | http client timeout ([go duration format](https://pkg.go.dev/time#Duration))    |
-| `PYROSCOPE_NUM_WORKERS`      | `5`                              | num of relay workers, pick based on the number of profile types                 |
-| `PYROSCOPE_FLUSH_ON_INVOKE`  | `false`                          | wait for all relay requests to be finished/flushed before next `Invocation` event is allowed |
+| env var                     | default                          | description                                                                                  |
+|-----------------------------|----------------------------------|----------------------------------------------------------------------------------------------|
+| `PYROSCOPE_REMOTE_ADDRESS`  | `https://ingest.pyroscope.cloud` | the pyroscope instance data will be relayed to                                               |
+| `PYROSCOPE_AUTH_TOKEN`      | `""`                             | authorization key (token authentication)                                                     |
+| `PYROSCOPE_SELF_PROFILING`  | `false`                          | whether to profile the extension itself or not                                               |
+| `PYROSCOPE_LOG_LEVEL`       | `info`                           | `error` or `info` or `debug` or `trace`                                                      |
+| `PYROSCOPE_TIMEOUT`         | `10s`                            | http client timeout ([go duration format](https://pkg.go.dev/time#Duration))                 |
+| `PYROSCOPE_NUM_WORKERS`     | `5`                              | num of relay workers, pick based on the number of profile types                              |
+| `PYROSCOPE_FLUSH_ON_INVOKE` | `false`                          | wait for all relay requests to be finished/flushed before next `Invocation` event is allowed |
+| `PYROSCOPE_HTTP_HEADERS`    | `{}`                             | extra http headers in json format, for example: {"X-Header": "Value"}                        |
+| `PYROSCOPE_SCOPE_ORGID`     | `""`                             | phlare tenant ID, passed as X-Scope-OrgID http header                                      |
+| `PYROSCOPE_BASIC_AUTH_USER`     | `""` | HTTP basic auth user |
+| `PYROSCOPE_BASIC_AUTH_PASSWORD`     | `""`  | HTTP basic auth password  |
 
 # How it works
 The profiler will run as normal, and periodically will send data to the relay server (the server running at `http://localhost:4040`).

go.mod

@@ -4,35 +4,33 @@ go 1.18
 
 require (
 	github.com/aws/aws-lambda-go v1.32.0
-	github.com/pyroscope-io/client v0.2.4-0.20220607180407-0ba26860ce5b
+	github.com/davecgh/go-spew v1.1.1
+	github.com/mgechev/revive v1.2.1
+	github.com/pyroscope-io/client v0.7.0
 	github.com/sirupsen/logrus v1.8.1
 	github.com/stretchr/testify v1.7.4
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
-	gotest.tools v2.2.0+incompatible
+	honnef.co/go/tools v0.3.2
 )
 
 require (
 	github.com/BurntSushi/toml v1.1.0 // indirect
 	github.com/chavacava/garif v0.0.0-20220316182200-5cad0b5181d4 // indirect
-	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fatih/color v1.13.0 // indirect
 	github.com/fatih/structtag v1.2.0 // indirect
-	github.com/google/go-cmp v0.5.7 // indirect
 	github.com/mattn/go-colorable v0.1.9 // indirect
 	github.com/mattn/go-isatty v0.0.14 // indirect
 	github.com/mattn/go-runewidth v0.0.9 // indirect
 	github.com/mgechev/dots v0.0.0-20210922191527-e955255bf517 // indirect
-	github.com/mgechev/revive v1.2.1 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
 	github.com/olekukonko/tablewriter v0.0.5 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/pyroscope-io/godeltaprof v0.1.0 // indirect
 	golang.org/x/exp/typeparams v0.0.0-20220218215828-6cf2b201936e // indirect
 	golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4 // indirect
 	golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e // indirect
 	golang.org/x/tools v0.1.11-0.20220513221640-090b14e8501f // indirect
-	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
-	honnef.co/go/tools v0.3.2 // indirect
 )

go.sum

@@ -11,8 +11,6 @@ github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
 github.com/fatih/structtag v1.2.0 h1:/OdNE99OxoI/PqaW/SuSK9uxxT3f/tcSZgon/ssNSx4=
 github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94=
-github.com/google/go-cmp v0.5.7 h1:81/ik6ipDQS2aGcBfIN5dHDB36BwrStyeAQquSYCV4o=
-github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8/DtOE=
 github.com/kr/pretty v0.2.1 h1:Fmg33tUaq4/8ym9TJN1x7sLJnHVwhP33CNkpYV/7rwI=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
@@ -37,8 +35,10 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/pyroscope-io/client v0.2.4-0.20220607180407-0ba26860ce5b h1:5YuLgMAQ7XbW8Kmccwv0C1vkl85SG7ItLC8AdMtCFOs=
-github.com/pyroscope-io/client v0.2.4-0.20220607180407-0ba26860ce5b/go.mod h1:zRdQXIGxy0H2QbKEkCmZBR6KOLLIFYLWsdzVI0MRm2E=
+github.com/pyroscope-io/client v0.7.0 h1:LWuuqPQ1oa6x7BnmUOuo/aGwdX85QGhWZUBYWWW3zdk=
+github.com/pyroscope-io/client v0.7.0/go.mod h1:4h21iOU4pUOq0prKyDlvYRL+SCKsBc5wKiEtV+rJGqU=
+github.com/pyroscope-io/godeltaprof v0.1.0 h1:UBqtjt0yZi4jTxqZmLAs34XG6ycS3vUTlhEUSq4NHLE=
+github.com/pyroscope-io/godeltaprof v0.1.0/go.mod h1:psMITXp90+8pFenXkKIpNhrfmI9saQnPbba27VIaiQE=
 github.com/sirupsen/logrus v1.8.1 h1:dJKuHgqk1NNQlqoA6BTlM1Wf9DOH3NBjQyu0h9+AZZE=
 github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -79,22 +79,17 @@ golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
-golang.org/x/tools v0.1.10 h1:QjFRCZxdOhBJ/UNgnBZLbNV13DlbnK0quyivTnXJM20=
 golang.org/x/tools v0.1.10/go.mod h1:Uh6Zz+xoGYZom868N8YTex3t7RhtHDBrE8Gzo9bV56E=
 golang.org/x/tools v0.1.11-0.20220513221640-090b14e8501f h1:OKYpQQVE3DKSc3r3zHVzq46vq5YH7x8xpR3/k9ixmUg=
 golang.org/x/tools v0.1.11-0.20220513221640-090b14e8501f/go.mod h1:SgwaegtQh8clINPpECJMqnxLv9I09HLqnW3RMqW0CA4=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-gotest.tools v2.2.0+incompatible h1:VsBPFP1AI068pPrMxtb/S8Zkgf9xEmTLJjfM+P5UIEo=
-gotest.tools v2.2.0+incompatible/go.mod h1:DsYFclhRJ6vuDpmuTbkuFWG+y2sxOXAzmJt81HFBacw=
 honnef.co/go/tools v0.3.2 h1:ytYb4rOqyp1TSa2EPvNVwtPQJctSELKaMyLfqNP4+34=
 honnef.co/go/tools v0.3.2/go.mod h1:jzwdWgg7Jdq75wlfblQxO4neNaFFSvgc1tD5Wv8U0Yw=

main.go

@@ -31,14 +31,19 @@ var (
 	// to where relay data to
 	remoteAddress = getEnvStrOr("PYROSCOPE_REMOTE_ADDRESS", "https://ingest.pyroscope.cloud")
 
-	authToken  = getEnvStrOr("PYROSCOPE_AUTH_TOKEN", "")
-	timeout    = getEnvDurationOr("PYROSCOPE_TIMEOUT", time.Second*10)
-	numWorkers = getEnvIntOr("PYROSCOPE_NUM_WORKERS", 5)
+	authToken         = getEnvStrOr("PYROSCOPE_AUTH_TOKEN", "")
+	basicAuthUser     = getEnvStrOr("PYROSCOPE_BASIC_AUTH_USER", "")
+	basicAuthPassword = getEnvStrOr("PYROSCOPE_BASIC_AUTH_PASSWORD", "")
+	scopeOrgID        = getEnvStrOr("PYROSCOPE_SCOPE_ORGID", "")
+	timeout           = getEnvDurationOr("PYROSCOPE_TIMEOUT", time.Second*10)
+	numWorkers        = getEnvIntOr("PYROSCOPE_NUM_WORKERS", 5)
 
 	// profile the extension?
 	selfProfiling = getEnvBool("PYROSCOPE_SELF_PROFILING")
 
 	flushOnInvoke = getEnvBool("PYROSCOPE_FLUSH_ON_INVOKE")
+
+	httpHeaders = getEnvStrOr("PYROSCOPE_HTTP_HEADERS", "")
 )
 
 func main() {
@@ -49,6 +54,10 @@ func main() {
 	remoteClient := relay.NewRemoteClient(logger, &relay.RemoteClientCfg{
 		Address:             remoteAddress,
 		AuthToken:           authToken,
+		BasicAuthUser:       basicAuthUser,
+		BasicAuthPassword:   basicAuthPassword,
+		ScopeOrgID:          scopeOrgID,
+		HTTPHeadersJSON:     httpHeaders,
 		Timeout:             timeout,
 		MaxIdleConnsPerHost: numWorkers,
 	})

relay/client.go

@@ -1,11 +1,13 @@
 package relay
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io/ioutil"
 	"net/http"
 	"net/url"
+	"path"
 	"time"
 
 	"github.com/sirupsen/logrus"
@@ -20,14 +22,19 @@ type RemoteClientCfg struct {
 	// Address refers to the remote address the request will be made to
 	Address             string
 	AuthToken           string
+	BasicAuthUser       string
+	BasicAuthPassword   string
+	ScopeOrgID          string
+	HTTPHeadersJSON     string
 	Timeout             time.Duration
 	MaxIdleConnsPerHost int
 }
 
 type RemoteClient struct {
-	config *RemoteClientCfg
-	client *http.Client
-	log    *logrus.Entry
+	config  *RemoteClientCfg
+	client  *http.Client
+	headers map[string]string
+	log     *logrus.Entry
 }
 
 func NewRemoteClient(log *logrus.Entry, config *RemoteClientCfg) *RemoteClient {
@@ -38,6 +45,13 @@ func NewRemoteClient(log *logrus.Entry, config *RemoteClientCfg) *RemoteClient {
 	if config.MaxIdleConnsPerHost == 0 {
 		config.MaxIdleConnsPerHost = 5
 	}
+	headers := make(map[string]string)
+	if config.HTTPHeadersJSON != "" {
+		err := json.Unmarshal([]byte(config.HTTPHeadersJSON), &headers)
+		if err != nil {
+			log.Error(fmt.Errorf("failed to parse headers json %w", err))
+		}
+	}
 	return &RemoteClient{
 		log:    log,
 		config: config,
@@ -56,6 +70,12 @@ func (r *RemoteClient) Send(req *http.Request) error {
 		defer req.Body.Close()
 	}
 	r.enhanceWithAuthToken(req)
+	if r.config.ScopeOrgID != "" {
+		req.Header.Set("X-Scope-OrgID", r.config.ScopeOrgID)
+	}
+	for k, v := range r.headers {
+		req.Header.Set(k, v)
+	}
 
 	host := r.config.Address
 
@@ -64,6 +84,8 @@ func (r *RemoteClient) Send(req *http.Request) error {
 	req.RequestURI = ""
 	req.URL.Host = u.Host
 	req.URL.Scheme = u.Scheme
+	req.URL.User = u.User
+	req.URL.Path = path.Join(u.Path, req.URL.Path)
 	req.Header.Set("X-Forwarded-Host", req.Header.Get("Host"))
 	req.Host = u.Host
 
@@ -91,5 +113,7 @@ func (r *RemoteClient) enhanceWithAuthToken(req *http.Request) {
 
 	if token != "" {
 		req.Header.Set("Authorization", "Bearer "+token)
+	} else if r.config.BasicAuthUser != "" && r.config.BasicAuthPassword != "" {
+		req.SetBasicAuth(r.config.BasicAuthUser, r.config.BasicAuthPassword)
 	}
 }