in progress

Author: prog-supdex

internal/commands/serve_rails.go internal/commands/serve.gointernal/commands/serve_rails.go renamed to internal/commands/serve.go

@@ -28,10 +28,10 @@ var logger = utils.NewCaddyRailsLogger()
 
 func init() {
 	caddycmd.RegisterCommand(caddycmd.Command{
-		Name:  "serve-rails",
+		Name:  "serve",
 		Short: "Runs an external server and sets up a reverse proxy to it",
 		Long: `
-The serve-rails command runs an external server specified as its argument and
+The serve command runs an external server specified as its argument and
 sets up a reverse proxy to forward requests to it.`,
 		CobraFunc: func(cmd *cobra.Command) {
 			cmd.Flags().String("pid-file", "", "Path to the PID file to control an existing process")

internal/utils/upstream_process.go

@@ -8,7 +8,10 @@ import (
 	"os/exec"
 	"os/signal"
 	"strconv"
+	"strings"
+	"sync"
 	"syscall"
+	"time"
 )
 
 const (
@@ -20,9 +23,12 @@ var logger = NewCaddyRailsLogger()
 
 type UpstreamProcess struct {
 	Started  chan struct{}
-	cmd      *exec.Cmd
+	Cmd      *exec.Cmd
 	SyncMode bool
 	PidFile  string
+	workers  []int
+	stopCh   chan struct{}
+	wg       sync.WaitGroup
 }
 
 func NewUpstreamProcess(command string, arg []string, syncMode bool, pidFile string) (*UpstreamProcess, error) {
@@ -34,97 +40,104 @@ func NewUpstreamProcess(command string, arg []string, syncMode bool, pidFile str
 
 	if command == "" && !fileExists(pidFile) {
 		logger.Error("For running an application, you must provide either an argument to the command serve-rails or ensure the presence of " + RailsExecutionFile)
-
 		return nil, fmt.Errorf("for running an application, you must provide either an argument to the command serve-rails or ensure the presence of %s file", RailsExecutionFile)
 	}
 
-	//logger.Info("Running a rails application by command: ", zap.String(command, strings.Join(arguments, " ")))
-
 	return &UpstreamProcess{
 		Started:  make(chan struct{}, 1),
-		cmd:      exec.Command(command, arguments...),
+		Cmd:      exec.Command(command, arguments...),
 		SyncMode: syncMode,
 		PidFile:  pidFile,
+		stopCh:   make(chan struct{}),
 	}, nil
 }
 
-func (p *UpstreamProcess) Run() (int, error) {
+func (p *UpstreamProcess) Run(testMode bool) (int, error) {
+	logger.Info("Starting Run method")
 	if p.PidFile != "" {
 		pid, err := p.readPidFile()
-
 		if err == nil {
-			p.cmd = &exec.Cmd{Process: &os.Process{Pid: pid}}
+			p.Cmd = &exec.Cmd{Process: &os.Process{Pid: pid}}
 			p.Started <- struct{}{}
 			if p.SyncMode {
-				return p.waitAndHandleExit()
+				select {}
 			}
-			go p.waitAndHandleExit()
 			return 0, nil
 		}
 	}
 
-	p.cmd.Stdout = os.Stdout
-	p.cmd.Stderr = os.Stderr
+	p.Cmd.Stdout = os.Stdout
+	p.Cmd.Stderr = os.Stderr
 
-	err := p.cmd.Start()
+	err := p.Cmd.Start()
 	if err != nil {
 		return 0, err
 	}
 
+	p.wg.Add(1)
 	p.Started <- struct{}{}
 
-	go p.handleSignals()
+	err = p.waitForPidFile()
+	if err != nil {
+		return 0, err
+	}
+
+	p.checkForNewWorkers()
+
+	// Signal handling
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		sig := <-sigs
+		logger.Info("Received signal", zap.String("signal", sig.String()))
+		p.terminateWorkers(sig)
+		p.Shutdown()
+		close(p.stopCh)
+	}()
 
 	if p.SyncMode {
-		return p.waitAndHandleExit()
+		<-p.stopCh
+		p.wg.Wait()
 	}
 
-	go p.waitAndHandleExit()
+	logger.Info("Run method finished successfully")
 	return 0, nil
 }
 
-func (p *UpstreamProcess) Stop() error {
-	if p.PidFile != "" {
-		pid, err := p.readPidFile()
+func (p *UpstreamProcess) Shutdown() error {
+	logger.Info("Starting Shutdown method")
+	if p.Cmd != nil && p.Cmd.Process != nil {
+		err := p.Cmd.Process.Signal(syscall.SIGTERM)
 		if err != nil {
 			return err
 		}
 
-		logger.Info("Stopping the rails application", zap.Int("pid", pid))
-
-		return p.signalProcess(pid, syscall.SIGTERM)
-	}
+		logger.Info("Waiting for main process to finish")
+		err = p.Cmd.Wait()
+		if err != nil {
+			logger.Error("Failed to wait for process", zap.Error(err))
+			return err
+		}
 
-	if p.cmd != nil && p.cmd.Process != nil {
-		return p.cmd.Process.Signal(syscall.SIGTERM)
+		p.wg.Done()
+		logger.Info("Main process finished")
+		return nil
 	}
 
 	return errors.New("process is not running or PID file is not specified")
 }
 
-func (p *UpstreamProcess) PhasedRestart(serverType string) error {
-	sig, err := determineSignal(serverType)
-	if err != nil {
-		return err
-	}
+func (p *UpstreamProcess) Stop() error {
+	if p.PidFile != "" {
+		pid, err := p.readPidFile()
+		if err != nil {
+			return err
+		}
 
-	pid, err := p.readPidFile()
-	if err != nil {
-		return err
+		return p.signalProcess(pid, syscall.SIGTERM)
 	}
 
-	logger.Info("Hot restarting the rails application", zap.Int("pid", pid))
-
-	return p.signalProcess(pid, sig)
-}
-
-func (p *UpstreamProcess) waitAndHandleExit() (int, error) {
-	err := p.cmd.Wait()
-	var exitErr *exec.ExitError
-	if errors.As(err, &exitErr) {
-		return exitErr.ExitCode(), nil
-	}
-	return 0, err
+	return errors.New("PidFile is not exists")
 }
 
 func (p *UpstreamProcess) readPidFile() (int, error) {
@@ -133,7 +146,7 @@ func (p *UpstreamProcess) readPidFile() (int, error) {
 		return 0, fmt.Errorf("failed to read PID file: %v", err)
 	}
 
-	pid, err := strconv.Atoi(string(pidData))
+	pid, err := strconv.Atoi(strings.TrimSpace(string(pidData)))
 	if err != nil {
 		return 0, fmt.Errorf("failed to convert PID to integer: %v", err)
 	}
@@ -149,19 +162,60 @@ func (p *UpstreamProcess) signalProcess(pid int, sig os.Signal) error {
 	return process.Signal(sig)
 }
 
-func (p *UpstreamProcess) signal(sig os.Signal) error {
-	return p.cmd.Process.Signal(sig)
+func (p *UpstreamProcess) waitForPidFile() error {
+	timeout := time.After(20 * time.Second)
+	tick := time.Tick(1 * time.Second)
+
+	for {
+		select {
+		case <-timeout:
+			return fmt.Errorf("timed out waiting for PID file")
+		case <-tick:
+			if fileExists(p.PidFile) {
+				logger.Info("PID file found")
+				return nil
+			}
+		}
+	}
 }
 
-func (p *UpstreamProcess) handleSignals() {
-	ch := make(chan os.Signal, 1)
-	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
+func (p *UpstreamProcess) checkForNewWorkers() {
+	output, err := exec.Command("pgrep", "-P", strconv.Itoa(p.Cmd.Process.Pid)).Output()
+	if err != nil {
+		logger.Warn("Failed to check for worker processes", zap.Error(err))
+		return
+	}
 
-	sig := <-ch
+	pidStrings := strings.Fields(string(output))
+	for _, pidStr := range pidStrings {
+		pid, err := strconv.Atoi(pidStr)
+		if err == nil {
+			logger.Info("Added worker with pid", zap.Int("pid", pid))
+			p.workers = append(p.workers, pid)
+			p.wg.Add(1) // Add worker to wait group
+		} else {
+			logger.Error("Error with worker pid", zap.Int("pid", pid), zap.Error(err))
+		}
+	}
+}
 
-	logger.Info("Relaying signal to upstream process", zap.String("signal", sig.String()))
+func (p *UpstreamProcess) terminateWorkers(sig os.Signal) {
+	for _, pid := range p.workers {
+		process, err := os.FindProcess(pid)
+		if err != nil {
+			logger.Warn("Failed to find worker process", zap.Int("pid", pid), zap.Error(err))
+			continue
+		}
+		err = process.Signal(sig)
+		if err != nil {
+			logger.Warn("Failed to signal worker process", zap.Int("pid", pid), zap.Error(err))
+		} else {
+			logger.Info("Signaled worker process", zap.Int("pid", pid), zap.String("signal", sig.String()))
+			p.wg.Done() // Mark worker as done
+		}
+	}
 
-	p.signal(sig)
+	p.workers = make([]int, 0)
 }
 
 func determineCommand(command string, arg []string) (string, []string) {
@@ -173,17 +227,6 @@ func determineCommand(command string, arg []string) (string, []string) {
 	return command, arg
 }
 
-func determineSignal(serverType string) (os.Signal, error) {
-	switch serverType {
-	case "puma":
-		return syscall.SIGUSR1, nil
-	case "unicorn":
-		return syscall.SIGUSR2, nil
-	default:
-		return nil, fmt.Errorf("unknown server type: %s", serverType)
-	}
-}
-
 func fileExists(filePath string) bool {
 	if _, err := os.Stat(filePath); err == nil {
 		return true