feat: allow NUM_WORKERS to be configurable (#12)

Author: korniltsev

README.md

@@ -30,6 +30,7 @@ For a complete list of variables check the section below.
 | `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 |
 
 # 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`).

main.go

@@ -31,8 +31,9 @@ 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)
+	authToken  = getEnvStrOr("PYROSCOPE_AUTH_TOKEN", "")
+	timeout    = getEnvDurationOr("PYROSCOPE_TIMEOUT", time.Second*10)
+	numWorkers = getEnvIntOr("PYROSCOPE_NUM_WORKERS", 5)
 
 	// profile the extension?
 	selfProfiling = getEnvBool("PYROSCOPE_SELF_PROFILING")
@@ -45,7 +46,7 @@ func main() {
 	// Init components
 	remoteClient := relay.NewRemoteClient(logger, &relay.RemoteClientCfg{Address: remoteAddress, AuthToken: authToken, Timeout: timeout})
 	// TODO(eh-am): a find a better default for num of workers
-	queue := relay.NewRemoteQueue(logger, &relay.RemoteQueueCfg{NumWorkers: 4}, remoteClient)
+	queue := relay.NewRemoteQueue(logger, &relay.RemoteQueueCfg{NumWorkers: numWorkers}, remoteClient)
 	ctrl := relay.NewController(logger, queue)
 	server := relay.NewServer(logger, &relay.ServerCfg{ServerAddress: "0.0.0.0:4040"}, ctrl.RelayRequest)
 
@@ -161,6 +162,7 @@ func getEnvStrOr(key string, fallback string) string {
 
 	return fallback
 }
+
 func getEnvBool(key string) bool {
 	k := os.Getenv(key)
 	v, err := strconv.ParseBool(k)
@@ -178,6 +180,7 @@ func getEnvDurationOr(key string, fallback time.Duration) time.Duration {
 	if ok && k != "" {
 		dur, err := time.ParseDuration(k)
 		if err != nil {
+			logrus.Warnf("invalid value for env var '%s': '%s' defaulting to '%s'", key, k, fallback)
 			return fallback
 		}
 
@@ -186,3 +189,19 @@ func getEnvDurationOr(key string, fallback time.Duration) time.Duration {
 
 	return fallback
 }
+
+func getEnvIntOr(key string, fallback int) int {
+	k, ok := os.LookupEnv(key)
+
+	// has an explicit value
+	if ok && k != "" {
+		val, err := strconv.Atoi(k)
+		if err != nil {
+			logrus.Warnf("invalid value for env var '%s': '%s' defaulting to '%d'", key, k, fallback)
+			return fallback
+		}
+		return val
+	}
+
+	return fallback
+}

relay/remotequeue.go

@@ -30,7 +30,7 @@ func NewRemoteQueue(log *logrus.Entry, config *RemoteQueueCfg, relayer Relayer)
 	// Setup defaults
 	if config.NumWorkers == 0 {
 		// TODO(eh-am): figure out a good default value?
-		config.NumWorkers = 4
+		config.NumWorkers = 5
 	}
 
 	return &RemoteQueue{