Skip to content

Commit

Permalink
🔥 Add OnPrefork Hooks so you can get the PID of the child process.
Browse files Browse the repository at this point in the history
  • Loading branch information
tohutohu committed Jul 10, 2022
1 parent 2a2337d commit fa0496a
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 0 deletions.
16 changes: 16 additions & 0 deletions hooks.go
Expand Up @@ -7,6 +7,7 @@ type OnGroupHandler = func(Group) error
type OnGroupNameHandler = OnGroupHandler
type OnListenHandler = func() error
type OnShutdownHandler = OnListenHandler
type OnForkHandler = func(int) error

type hooks struct {
// Embed app
Expand All @@ -19,6 +20,7 @@ type hooks struct {
onGroupName []OnGroupNameHandler
onListen []OnListenHandler
onShutdown []OnShutdownHandler
onFork []OnForkHandler
}

func newHooks(app *App) *hooks {
Expand All @@ -30,6 +32,7 @@ func newHooks(app *App) *hooks {
onName: make([]OnNameHandler, 0),
onListen: make([]OnListenHandler, 0),
onShutdown: make([]OnShutdownHandler, 0),
onFork: make([]OnForkHandler, 0),
}
}

Expand Down Expand Up @@ -83,6 +86,13 @@ func (h *hooks) OnShutdown(handler ...OnShutdownHandler) {
h.app.mutex.Unlock()
}

// OnFork is a hook to execute user function after fork process.
func (h *hooks) OnFork(handler ...OnForkHandler) {
h.app.mutex.Lock()
h.onFork = append(h.onFork, handler...)
h.app.mutex.Unlock()
}

func (h *hooks) executeOnRouteHooks(route Route) error {
for _, v := range h.onRoute {
if err := v(route); err != nil {
Expand Down Expand Up @@ -138,3 +148,9 @@ func (h *hooks) executeOnShutdownHooks() {
_ = v()
}
}

func (h *hooks) executeOnForkHooks(pid int) {
for _, v := range h.onFork {
_ = v(pid)
}
}
20 changes: 20 additions & 0 deletions hooks_test.go
Expand Up @@ -176,3 +176,23 @@ func Test_Hook_OnListen(t *testing.T) {

utils.AssertEqual(t, "ready", buf.String())
}

func Test_Hook_OnHook(t *testing.T) {
// Reset test var
testPreforkMaster = true
testOnPrefork = true

app := New()

go func() {
time.Sleep(1000 * time.Millisecond)
utils.AssertEqual(t, nil, app.Shutdown())
}()

app.Hooks().OnFork(func(pid int) error {
utils.AssertEqual(t, 1, pid)
return nil
})

utils.AssertEqual(t, nil, app.prefork(NetworkTCP4, ":3000", nil))
}
12 changes: 12 additions & 0 deletions prefork.go
Expand Up @@ -20,6 +20,7 @@ const (
)

var testPreforkMaster = false
var testOnPrefork = false

// IsChild determines if the current process is a child of Prefork
func IsChild() bool {
Expand Down Expand Up @@ -102,6 +103,15 @@ func (app *App) prefork(network, addr string, tlsConfig *tls.Config) (err error)
childs[pid] = cmd
pids = append(pids, strconv.Itoa(pid))

// execute fork hook
if app.hooks != nil {
if testOnPrefork {
app.hooks.executeOnForkHooks(dummyPid)
} else {
app.hooks.executeOnForkHooks(pid)
}
}

// notify master if child crashes
go func() {
channel <- child{pid, cmd.Wait()}
Expand Down Expand Up @@ -146,3 +156,5 @@ func dummyCmd() *exec.Cmd {
}
return exec.Command(dummyChildCmd, "version")
}

var dummyPid = 1

0 comments on commit fa0496a

Please sign in to comment.