Make the prefork mode more robust (#755)

* Make the prefork mode more robust

The main process will exit if one of the prefork child processes doesn't complete successfully under the current prefork mode, so it ought to make sure that all child processes run independently and the main process will only exit after all child processes are finished.

* Start over those failed child processes automatically

* Kill all child processes before main process exits

* Remove redundant code

* Add configurable threshold of starting over child processes

* Return a error of RecoverThreshold

* Resolved requested changes

* Add logs

* Resolve requested changes

1 files changed, 102 insertions, 15 deletions

diff --git a/prefork/prefork.go b/prefork/prefork.go
index b7a1470..0b31ef5 100644
--- a/prefork/prefork.go
+++ b/prefork/prefork.go
@@ -1,7 +1,9 @@
 package prefork
 
 import (
+	"errors"
 	"flag"
+	"log"
 	"net"
 	"os"
 	"os/exec"
@@ -11,8 +13,23 @@ import (
 	"github.com/valyala/fasthttp/reuseport"
 )
 
-const preforkChildFlag = "-prefork-child"
-const defaultNetwork = "tcp4"
+const (
+	preforkChildFlag = "-prefork-child"
+	defaultNetwork   = "tcp4"
+)
+
+var (
+	defaultLogger = Logger(log.New(os.Stderr, "", log.LstdFlags))
+	// ErrOverRecovery is returned when the times of starting over child prefork processes exceed
+	// the threshold.
+	ErrOverRecovery = errors.New("exceeding the value of RecoverThreshold")
+)
+
+// Logger is used for logging formatted messages.
+type Logger interface {
+	// Printf must have the same semantics as log.Printf.
+	Printf(format string, args ...interface{})
+}
 
 // Prefork implements fasthttp server prefork
 //
@@ -34,6 +51,13 @@ type Prefork struct {
 	// It's disabled by default
 	Reuseport bool
 
+	// Child prefork processes may exit with failure and will be started over until the times reach
+	// the value of RecoverThreshold, then it will return and terminate the server.
+	RecoverThreshold int
+
+	// By default standard logger from log package is used.
+	Logger Logger
+
 	ServeFunc         func(ln net.Listener) error
 	ServeTLSFunc      func(ln net.Listener, certFile, keyFile string) error
 	ServeTLSEmbedFunc func(ln net.Listener, certData, keyData []byte) error
@@ -63,12 +87,21 @@ func IsChild() bool {
 func New(s *fasthttp.Server) *Prefork {
 	return &Prefork{
 		Network:           defaultNetwork,
+		RecoverThreshold:  runtime.GOMAXPROCS(0) / 2,
+		Logger:            s.Logger,
 		ServeFunc:         s.Serve,
 		ServeTLSFunc:      s.ServeTLS,
 		ServeTLSEmbedFunc: s.ServeTLSEmbed,
 	}
 }
 
+func (p *Prefork) logger() Logger {
+	if p.Logger != nil {
+		return p.Logger
+	}
+	return defaultLogger
+}
+
 func (p *Prefork) listen(addr string) (net.Listener, error) {
 	runtime.GOMAXPROCS(1)
 
@@ -110,30 +143,84 @@ func (p *Prefork) setTCPListenerFiles(addr string) error {
 	return nil
 }
 
-func (p *Prefork) prefork(addr string) error {
-	chErr := make(chan error, 1)
+func (p *Prefork) doCommand() (*exec.Cmd, error) {
+	/* #nosec G204 */
+	cmd := exec.Command(os.Args[0], append(os.Args[1:], preforkChildFlag)...)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	cmd.ExtraFiles = p.files
+	return cmd, cmd.Start()
+}
 
+func (p *Prefork) prefork(addr string) (err error) {
 	if !p.Reuseport {
-		if err := p.setTCPListenerFiles(addr); err != nil {
-			return err
+		if err = p.setTCPListenerFiles(addr); err != nil {
+			return
+		}
+
+		// defer for closing the net.Listener opened by setTCPListenerFiles.
+		defer func() {
+			e := p.ln.Close()
+			if err == nil {
+				err = e
+			}
+		}()
+	}
+
+	type procSig struct {
+		pid int
+		err error
+	}
+
+	goMaxProcs := runtime.GOMAXPROCS(0)
+	sigCh := make(chan procSig, goMaxProcs)
+	childProcs := make(map[int]*exec.Cmd)
+
+	defer func() {
+		for _, proc := range childProcs {
+			_ = proc.Process.Kill()
+		}
+	}()
+
+	for i := 0; i < goMaxProcs; i++ {
+		var cmd *exec.Cmd
+		if cmd, err = p.doCommand(); err != nil {
+			p.logger().Printf("failed to start a child prefork process, error: %v\n", err)
+			return
 		}
 
-		defer p.ln.Close()
+		childProcs[cmd.Process.Pid] = cmd
+		go func() {
+			sigCh <- procSig{cmd.Process.Pid, cmd.Wait()}
+		}()
 	}
 
-	for i := 0; i < runtime.GOMAXPROCS(0); i++ {
-		/* #nosec G204 */
-		cmd := exec.Command(os.Args[0], append(os.Args[1:], preforkChildFlag)...)
-		cmd.Stdout = os.Stdout
-		cmd.Stderr = os.Stderr
-		cmd.ExtraFiles = p.files
+	var exitedProcs int
+	for sig := range sigCh {
+		delete(childProcs, sig.pid)
+
+		p.logger().Printf("one of the child prefork processes exited with "+
+			"error: %v", sig.err)
 
+		if exitedProcs++; exitedProcs > p.RecoverThreshold {
+			p.logger().Printf("child prefork processes exit too many times, "+
+				"which exceeds the value of RecoverThreshold(%d), "+
+				"exiting the master process.\n", exitedProcs)
+			err = ErrOverRecovery
+			break
+		}
+
+		var cmd *exec.Cmd
+		if cmd, err = p.doCommand(); err != nil {
+			break
+		}
+		childProcs[cmd.Process.Pid] = cmd
 		go func() {
-			chErr <- cmd.Run()
+			sigCh <- procSig{cmd.Process.Pid, cmd.Wait()}
 		}()
 	}
 
-	return <-chErr
+	return
 }
 
 // ListenAndServe serves HTTP requests from the given TCP addr