Add prefork utility

prefork/prefork.go

@@ -0,0 +1,179 @@
+package prefork
+
+import (
+	"flag"
+	"net"
+	"os"
+	"os/exec"
+	"runtime"
+
+	"github.com/valyala/fasthttp"
+	"github.com/valyala/fasthttp/reuseport"
+)
+
+const preforkChildFlag = "prefork-child"
+const defaultNetwork = "tcp4"
+
+var Child bool
+
+// Prefork implements fasthttp server prefork
+//
+// Preforks master process (with all cores) between several child processes
+// increases performance significantly, because Go doesn't have to share
+// and manage memory between cores
+//
+// WARNING: Does not recommended for servers with in-memory cache,
+// because the cache will be duplicated in each process
+type Prefork struct {
+	Addr string
+
+	// The network must be "tcp", "tcp4" or "tcp6".
+	//
+	// By default is "tcp4"
+	Network string
+
+	// Flag to use a listener with reuseport, if not a File Listener will be used
+	//
+	// It's disabled by default
+	Reuseport bool
+
+	ServeFunc         func(ln net.Listener) error
+	ServeTLSFunc      func(ln net.Listener, certFile, keyFile string) error
+	ServeTLSEmbedFunc func(ln net.Listener, certData, keyData []byte) error
+
+	ln    net.Listener
+	files []*os.File
+}
+
+func init() { // nolint:gochecknoinits
+	flag.BoolVar(&Child, preforkChildFlag, false, "is child proc")
+}
+
+// New wraps the fasthttp server to run with prefork processes
+func New(s *fasthttp.Server) *Prefork {
+	return &Prefork{
+		Network:           defaultNetwork,
+		ServeFunc:         s.Serve,
+		ServeTLSFunc:      s.ServeTLS,
+		ServeTLSEmbedFunc: s.ServeTLSEmbed,
+	}
+}
+
+func (p *Prefork) listen(addr string) (net.Listener, error) {
+	p.Addr = addr
+
+	runtime.GOMAXPROCS(1)
+
+	if p.Network == "" {
+		p.Network = defaultNetwork
+	}
+
+	if p.Reuseport {
+		return reuseport.Listen(p.Network, p.Addr)
+	}
+
+	return net.FileListener(os.NewFile(3, ""))
+}
+
+func (p *Prefork) setTCPListenerFiles(addr string) error {
+	p.Addr = addr
+
+	tcpAddr, err := net.ResolveTCPAddr(p.Network, p.Addr)
+	if err != nil {
+		return err
+	}
+
+	tcplistener, err := net.ListenTCP(p.Network, tcpAddr)
+	if err != nil {
+		return err
+	}
+
+	p.ln = tcplistener
+
+	fl, err := tcplistener.File()
+	if err != nil {
+		return err
+	}
+
+	p.files = []*os.File{fl}
+
+	return nil
+}
+
+func (p *Prefork) prefork(addr string) error {
+	strCmd := os.Args[0]
+	chErr := make(chan error, 1)
+
+	if !p.Reuseport {
+		if err := p.setTCPListenerFiles(addr); err != nil {
+			return err
+		}
+
+		defer p.ln.Close()
+	}
+
+	for i := 0; i < runtime.NumCPU(); i++ {
+		cmd := exec.Command(strCmd, append(os.Args[1:], "-"+preforkChildFlag)...)
+		cmd.Stdout = os.Stdout
+		cmd.Stderr = os.Stderr
+		cmd.ExtraFiles = p.files
+
+		go func() {
+			chErr <- cmd.Run()
+		}()
+	}
+
+	return <-chErr
+}
+
+// ListenAndServe serves HTTP requests from the given TCP addr
+func (p *Prefork) ListenAndServe(addr string) error {
+	if Child {
+		ln, err := p.listen(addr)
+		if err != nil {
+			return err
+		}
+
+		p.ln = ln
+
+		return p.ServeFunc(ln)
+	}
+
+	return p.prefork(addr)
+}
+
+// ListenAndServeTLS serves HTTPS requests from the given TCP addr
+//
+// certFile and keyFile are paths to TLS certificate and key files.
+func (p *Prefork) ListenAndServeTLS(addr, certKey, certFile string) error {
+	if Child {
+		ln, err := p.listen(addr)
+		if err != nil {
+			return err
+		}
+
+		p.ln = ln
+
+		return p.ServeTLSFunc(ln, certFile, certKey)
+	}
+
+	return p.prefork(addr)
+}
+
+// ListenAndServeTLSEmbed serves HTTPS requests from the given TCP addr
+//
+// certData and keyData must contain valid TLS certificate and key data.
+func (p *Prefork) ListenAndServeTLSEmbed(addr string, certData, keyData []byte) error {
+	if Child {
+		ln, err := p.listen(addr)
+		if err != nil {
+			return err
+		}
+
+		p.ln = ln
+
+		return p.ServeTLSEmbedFunc(ln, certData, keyData)
+	}
+
+	return p.prefork(addr)
+}

prefork/prefork_test.go

@@ -0,0 +1,172 @@
+package prefork
+
+import (
+	"fmt"
+	"math/rand"
+	"net"
+	"reflect"
+	"runtime"
+	"testing"
+
+	"github.com/valyala/fasthttp"
+)
+
+func getAddr() string {
+	Child = true
+
+	return fmt.Sprintf(":%d", rand.Intn(9000-3000)+3000)
+}
+
+func Test_New(t *testing.T) {
+	s := &fasthttp.Server{}
+	p := New(s)
+
+	if p.Network != defaultNetwork {
+		t.Errorf("Prefork.Netork == %s, want %s", p.Network, defaultNetwork)
+	}
+
+	if reflect.ValueOf(p.ServeFunc).Pointer() != reflect.ValueOf(s.Serve).Pointer() {
+		t.Errorf("Prefork.ServeFunc == %p, want %p", p.ServeFunc, s.Serve)
+	}
+
+	if reflect.ValueOf(p.ServeTLSFunc).Pointer() != reflect.ValueOf(s.ServeTLS).Pointer() {
+		t.Errorf("Prefork.ServeTLSFunc == %p, want %p", p.ServeTLSFunc, s.ServeTLS)
+	}
+
+	if reflect.ValueOf(p.ServeTLSEmbedFunc).Pointer() != reflect.ValueOf(s.ServeTLSEmbed).Pointer() {
+		t.Errorf("Prefork.ServeTLSFunc == %p, want %p", p.ServeTLSEmbedFunc, s.ServeTLSEmbed)
+	}
+}
+
+func Test_listen(t *testing.T) {
+	p := &Prefork{
+		Reuseport: true,
+	}
+	addr := getAddr()
+
+	ln, err := p.listen(addr)
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	ln.Close()
+
+	if p.Addr != addr {
+		t.Errorf("Prefork.Addr == %s, want %s", p.Addr, addr)
+	}
+
+	if p.Network != defaultNetwork {
+		t.Errorf("Prefork.Network == %s, want %s", p.Network, defaultNetwork)
+	}
+
+	procs := runtime.GOMAXPROCS(0)
+	if procs != 1 {
+		t.Errorf("GOMAXPROCS == %d, want %d", procs, 1)
+	}
+}
+
+func Test_setTCPListenerFiles(t *testing.T) {
+	p := &Prefork{
+		Network: defaultNetwork,
+	}
+	addr := getAddr()
+
+	err := p.setTCPListenerFiles(addr)
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if p.Addr != addr {
+		t.Errorf("Prefork.Addr == %s, want %s", p.Addr, addr)
+	}
+
+	if p.ln == nil {
+		t.Fatal("Prefork.ln is nil")
+	}
+
+	p.ln.Close()
+
+	if len(p.files) != 1 {
+		t.Errorf("Prefork.files == %d, want %d", len(p.files), 1)
+	}
+}
+
+func Test_ListenAndServe(t *testing.T) {
+	s := &fasthttp.Server{}
+	p := New(s)
+	p.Reuseport = true
+	p.ServeFunc = func(ln net.Listener) error {
+		return nil
+	}
+
+	addr := getAddr()
+
+	err := p.ListenAndServe(addr)
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	p.ln.Close()
+
+	if p.Addr != addr {
+		t.Errorf("Prefork.Addr == %s, want %s", p.Addr, addr)
+	}
+
+	if p.ln == nil {
+		t.Error("Prefork.ln is nil")
+	}
+}
+
+func Test_ListenAndServeTLS(t *testing.T) {
+	s := &fasthttp.Server{}
+	p := New(s)
+	p.Reuseport = true
+	p.ServeTLSFunc = func(ln net.Listener, certFile, keyFile string) error {
+		return nil
+	}
+
+	addr := getAddr()
+
+	err := p.ListenAndServeTLS(addr, "./key", "./cert")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	p.ln.Close()
+
+	if p.Addr != addr {
+		t.Errorf("Prefork.Addr == %s, want %s", p.Addr, addr)
+	}
+
+	if p.ln == nil {
+		t.Error("Prefork.ln is nil")
+	}
+}
+
+func Test_ListenAndServeTLSEmbed(t *testing.T) {
+	s := &fasthttp.Server{}
+	p := New(s)
+	p.Reuseport = true
+	p.ServeTLSEmbedFunc = func(ln net.Listener, certData, keyData []byte) error {
+		return nil
+	}
+
+	addr := getAddr()
+
+	err := p.ListenAndServeTLSEmbed(addr, []byte("key"), []byte("cert"))
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	p.ln.Close()
+
+	if p.Addr != addr {
+		t.Errorf("Prefork.Addr == %s, want %s", p.Addr, addr)
+	}
+
+	if p.ln == nil {
+		t.Error("Prefork.ln is nil")
+	}
+}