feat: Add PC_REPLICA_NUM variable to each process

F1bonacc1 · Dec 20, 2024 · d7e0380 · d7e0380
1 parent c0e29e4
commit d7e0380
Show file tree

Hide file tree

Showing 9 changed files with 105 additions and 40 deletions.
diff --git a/process-compose.yaml b/process-compose.yaml
@@ -98,8 +98,13 @@ processes:
       failure_threshold: 3
 
   server:
-    command: "python3 -m http.server 4040"
+    command: "python3 -m http.server 404{{.PC_REPLICA_NUM}}"
     is_tty: true
+    readiness_probe:
+      http_get:
+        host: "127.0.0.1"
+        port: "404{{.PC_REPLICA_NUM}}"
+        scheme: "http"
 
   kcalc:
     command: "kcalc"

diff --git a/src/app/project_runner.go b/src/app/project_runner.go
@@ -2,12 +2,14 @@ package app
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"github.com/f1bonacc1/process-compose/src/config"
 	"github.com/f1bonacc1/process-compose/src/health"
 	"github.com/f1bonacc1/process-compose/src/loader"
 	"github.com/f1bonacc1/process-compose/src/pclog"
+	"github.com/f1bonacc1/process-compose/src/templater"
 	"github.com/f1bonacc1/process-compose/src/types"
 	"os"
 	"os/user"
@@ -669,10 +671,19 @@ func (p *ProjectRunner) getCurrentReplicaCount(name string) int {
 
 func (p *ProjectRunner) scaleUpProcess(proc types.ProcessConfig, toAdd, scale, origScale int) {
 	for i := 0; i < toAdd; i++ {
-		proc.ReplicaNum = origScale + i
-		proc.Replicas = scale
-		proc.ReplicaName = proc.CalculateReplicaName()
-		p.addProcessAndRun(proc)
+		var procFromConf types.ProcessConfig
+		err := json.Unmarshal([]byte(proc.OriginalConfig), &procFromConf)
+		if err != nil {
+			log.Err(err).Msgf("failed to unmarshal config for %s", proc.Name)
+			return
+		}
+		procFromConf.ReplicaNum = origScale + i
+		procFromConf.Replicas = scale
+		procFromConf.ReplicaName = procFromConf.CalculateReplicaName()
+		tpl := templater.New(p.project.Vars)
+		tpl.RenderProcess(&procFromConf)
+		procFromConf.AssignProcessExecutableAndArgs(p.project.ShellConfig, p.project.GetElevatedShellArg())
+		p.addProcessAndRun(procFromConf)
 	}
 }
 
@@ -735,11 +746,11 @@ func (p *ProjectRunner) renameProcess(name string, newName string) {
 		state.Name = newName
 		p.processStates[newName] = state
 	}
-	config, ok := p.project.Processes[name]
+	procConf, ok := p.project.Processes[name]
 	if ok {
 		delete(p.project.Processes, name)
-		config.ReplicaName = newName
-		p.project.Processes[newName] = config
+		procConf.ReplicaName = newName
+		p.project.Processes[newName] = procConf
 	}
 }
 func (p *ProjectRunner) removeProcessLogs(name string) *pclog.ProcessLogBuffer {

diff --git a/src/loader/mutators.go b/src/loader/mutators.go
@@ -3,7 +3,6 @@ package loader
 import (
 	"fmt"
 	"github.com/f1bonacc1/process-compose/src/command"
-	"github.com/f1bonacc1/process-compose/src/health"
 	"github.com/f1bonacc1/process-compose/src/templater"
 	"github.com/f1bonacc1/process-compose/src/types"
 	"github.com/rs/zerolog/log"
@@ -115,10 +114,7 @@ func cloneReplicas(p *types.Project) {
 }
 
 func assignExecutableAndArgs(p *types.Project) {
-	elevatedShellArg := p.ShellConfig.ElevatedShellArg
-	if p.IsTuiDisabled {
-		elevatedShellArg = ""
-	}
+	elevatedShellArg := p.GetElevatedShellArg()
 	for name, proc := range p.Processes {
 		proc.AssignProcessExecutableAndArgs(p.ShellConfig, elevatedShellArg)
 
@@ -129,15 +125,7 @@ func assignExecutableAndArgs(p *types.Project) {
 func renderTemplates(p *types.Project) error {
 	tpl := templater.New(p.Vars)
 	for name, proc := range p.Processes {
-		if len(p.Vars) == 0 && len(proc.Vars) == 0 {
-			continue
-		}
-		proc.Command = tpl.RenderWithExtraVars(proc.Command, proc.Vars)
-		proc.WorkingDir = tpl.RenderWithExtraVars(proc.WorkingDir, proc.Vars)
-		proc.LogLocation = tpl.RenderWithExtraVars(proc.LogLocation, proc.Vars)
-		proc.Description = tpl.RenderWithExtraVars(proc.Description, proc.Vars)
-		renderProbe(proc.ReadinessProbe, tpl, proc.Vars)
-		renderProbe(proc.LivenessProbe, tpl, proc.Vars)
+		tpl.RenderProcess(&proc)
 
 		if tpl.GetError() != nil {
 			return fmt.Errorf("error rendering template for process %s: %w", name, tpl.GetError())
@@ -146,19 +134,3 @@ func renderTemplates(p *types.Project) error {
 	}
 	return nil
 }
-
-func renderProbe(probe *health.Probe, tpl *templater.Templater, vars types.Vars) {
-	if probe == nil {
-		return
-	}
-
-	if probe.Exec != nil {
-		probe.Exec.Command = tpl.RenderWithExtraVars(probe.Exec.Command, vars)
-	} else if probe.HttpGet != nil {
-		probe.HttpGet.Path = tpl.RenderWithExtraVars(probe.HttpGet.Path, vars)
-		probe.HttpGet.Host = tpl.RenderWithExtraVars(probe.HttpGet.Host, vars)
-		probe.HttpGet.Scheme = tpl.RenderWithExtraVars(probe.HttpGet.Scheme, vars)
-		probe.HttpGet.Port = tpl.RenderWithExtraVars(probe.HttpGet.Port, vars)
-	}
-	probe.ValidateAndSetDefaults()
-}
diff --git a/src/loader/mutators_test.go b/src/loader/mutators_test.go
@@ -449,7 +449,7 @@ func Test_renderTemplates(t *testing.T) {
 			for _, p := range tt.args.p.Processes {
 				switch tt.name {
 				case "no vars":
-					expected := "echo {{ .TEST }}"
+					expected := "echo <no value>"
 					compareStrings(t, expected, p.ReadinessProbe.Exec.Command, "readiness probe command")
 					compareStrings(t, expected, p.LivenessProbe.Exec.Command, "liveness probe command")
 					compareStrings(t, expected, p.Command, "process command")
@@ -480,6 +480,7 @@ func Test_renderTemplates(t *testing.T) {
 }
 
 func compareStrings(t *testing.T, expected, actual, scope string) {
+	t.Helper()
 	if expected != actual {
 		t.Errorf("Expected %s '%s' to be '%s'", scope, expected, actual)
 	}

diff --git a/src/templater/templater.go b/src/templater/templater.go
@@ -2,7 +2,10 @@ package templater
 
 import (
 	"bytes"
+	"encoding/json"
+	"github.com/f1bonacc1/process-compose/src/health"
 	"github.com/f1bonacc1/process-compose/src/types"
+	"github.com/rs/zerolog/log"
 	"maps"
 	"text/template"
 )
@@ -16,6 +19,40 @@ func New(vars types.Vars) *Templater {
 	return &Templater{vars: vars}
 }
 
+func (t *Templater) RenderProcess(proc *types.ProcessConfig) {
+	if proc.Vars == nil {
+		proc.Vars = make(types.Vars)
+	}
+	procConf, err := json.Marshal(proc)
+	if err != nil {
+		log.Error().Err(err).Msg("Failed to marshal process config")
+	}
+	proc.OriginalConfig = string(procConf)
+	proc.Vars["PC_REPLICA_NUM"] = proc.ReplicaNum
+	proc.Command = t.RenderWithExtraVars(proc.Command, proc.Vars)
+	proc.WorkingDir = t.RenderWithExtraVars(proc.WorkingDir, proc.Vars)
+	proc.LogLocation = t.RenderWithExtraVars(proc.LogLocation, proc.Vars)
+	proc.Description = t.RenderWithExtraVars(proc.Description, proc.Vars)
+	t.renderProbe(proc.ReadinessProbe, proc)
+	t.renderProbe(proc.LivenessProbe, proc)
+}
+
+func (t *Templater) renderProbe(probe *health.Probe, procConf *types.ProcessConfig) {
+	if probe == nil {
+		return
+	}
+
+	if probe.Exec != nil {
+		probe.Exec.Command = t.RenderWithExtraVars(probe.Exec.Command, procConf.Vars)
+	} else if probe.HttpGet != nil {
+		probe.HttpGet.Path = t.RenderWithExtraVars(probe.HttpGet.Path, procConf.Vars)
+		probe.HttpGet.Host = t.RenderWithExtraVars(probe.HttpGet.Host, procConf.Vars)
+		probe.HttpGet.Scheme = t.RenderWithExtraVars(probe.HttpGet.Scheme, procConf.Vars)
+		probe.HttpGet.Port = t.RenderWithExtraVars(probe.HttpGet.Port, procConf.Vars)
+	}
+	probe.ValidateAndSetDefaults()
+}
+
 func (t *Templater) Render(str string) string {
 	return t.render(str, nil)
 }

diff --git a/src/templater/templater_test.go b/src/templater/templater_test.go
@@ -48,6 +48,20 @@ func TestTemplater_RenderWithExtraVars(t *testing.T) {
 			t.Errorf("Expected %s but got %s", expected, result)
 		}
 	})
+	t.Run("Rendering with proc conf", func(t *testing.T) {
+		vars := types.Vars{"Name": "Alice"}
+
+		procConf := &types.ProcessConfig{
+			ReplicaNum: 3,
+			Command:    "Name: {{.Name}}, Replica: {{.PC_REPLICA_NUM}}",
+		}
+		templater := New(vars)
+		templater.RenderProcess(procConf)
+		expected := "Name: Alice, Replica: 3"
+		if procConf.Command != expected {
+			t.Errorf("Expected %s but got %s", expected, procConf.Command)
+		}
+	})
 }
 
 func TestTemplater_GetError(t *testing.T) {

diff --git a/src/types/process.go b/src/types/process.go
@@ -44,6 +44,7 @@ type ProcessConfig struct {
 	IsTty             bool                   `yaml:"is_tty"`
 	IsElevated        bool                   `yaml:"is_elevated"`
 	LaunchTimeout     int                    `yaml:"launch_timeout_seconds"`
+	OriginalConfig    string
 	ReplicaNum        int
 	ReplicaName       string
 	Executable        string

diff --git a/src/types/project.go b/src/types/project.go
@@ -56,6 +56,14 @@ func (p *Project) GetLexicographicProcessNames() ([]string, error) {
 	return names, nil
 }
 
+func (p *Project) GetElevatedShellArg() string {
+	elevatedShellArg := p.ShellConfig.ElevatedShellArg
+	if p.IsTuiDisabled {
+		elevatedShellArg = ""
+	}
+	return elevatedShellArg
+}
+
 func (p *Project) GetProcesses(names ...string) ([]ProcessConfig, error) {
 	processes := []ProcessConfig{}
 	if len(names) == 0 {

diff --git a/www/docs/configuration.md b/www/docs/configuration.md
@@ -202,7 +202,23 @@ processes:
         command: '[ $(docker inspect -f {{ "{{.State.Running}}" }} nginx_test) = true ]'
 ```
 
-> :bulb: For backward compatibility, if neither global nor local variables exist in `process-compose.yaml` the template engine won't run.
+### Auto Inserted Variables
+
+Process Compose will insert `PC_REPLICA_NUM` variable that will represent the replica number of the process. This will allow to conveniently scale processes using the following example configuration:
+
+```yaml hl_lines="3 8"
+processes:  
+  server:
+    command: "python3 -m http.server 404{{.PC_REPLICA_NUM}}"
+    is_tty: true
+    readiness_probe:
+      http_get:
+        host: "127.0.0.1"
+        port: "404{{.PC_REPLICA_NUM}}"
+        scheme: "http"
+```
+
+
 
 ## Specify which configuration files to use