Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

πŸ”₯ Add OnPrefork Hooks so you can get the PID of the child process. #1974

Merged
merged 1 commit into from Aug 1, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
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