diff --git a/daemon.go b/daemon.go index 2e63d02..5adbb26 100644 --- a/daemon.go +++ b/daemon.go @@ -179,10 +179,9 @@ func (d *Daemon) Stop() error { d.health.SetReady(false) } - // Auto-unregister after the process is no longer serving traffic. - if d.opts.Registry != nil { - if err := d.opts.Registry.Unregister(d.opts.RegistryEntry.Code, d.opts.RegistryEntry.Daemon); err != nil { - errs = append(errs, coreerr.E("Daemon.Stop", "registry", err)) + if d.health != nil { + if err := d.health.Stop(shutdownCtx); err != nil { + errs = append(errs, coreerr.E("Daemon.Stop", "health server", err)) } } @@ -192,9 +191,11 @@ func (d *Daemon) Stop() error { } } - if d.health != nil { - if err := d.health.Stop(shutdownCtx); err != nil { - errs = append(errs, coreerr.E("Daemon.Stop", "health server", err)) + // Auto-unregister after the daemon has stopped serving traffic and + // relinquished its PID file. + if d.opts.Registry != nil { + if err := d.opts.Registry.Unregister(d.opts.RegistryEntry.Code, d.opts.RegistryEntry.Daemon); err != nil { + errs = append(errs, coreerr.E("Daemon.Stop", "registry", err)) } } diff --git a/daemon_test.go b/daemon_test.go index f32a0e7..2c12333 100644 --- a/daemon_test.go +++ b/daemon_test.go @@ -109,7 +109,7 @@ func TestDaemon_StopMarksNotReadyBeforeShutdownCompletes(t *testing.T) { } } -func TestDaemon_StopUnregistersBeforeHealthShutdownCompletes(t *testing.T) { +func TestDaemon_StopUnregistersAfterHealthShutdownCompletes(t *testing.T) { blockCheck := make(chan struct{}) checkEntered := make(chan struct{}) var once sync.Once @@ -165,10 +165,8 @@ func TestDaemon_StopUnregistersBeforeHealthShutdownCompletes(t *testing.T) { return !d.Ready() }, 500*time.Millisecond, 10*time.Millisecond, "daemon should become not ready before shutdown completes") - require.Eventually(t, func() bool { - _, ok := reg.Get("test-app", "serve") - return !ok - }, 500*time.Millisecond, 10*time.Millisecond, "daemon should unregister before health shutdown completes") + _, ok := reg.Get("test-app", "serve") + assert.True(t, ok, "daemon should remain registered until health shutdown completes") select { case err := <-stopDone: @@ -185,6 +183,11 @@ func TestDaemon_StopUnregistersBeforeHealthShutdownCompletes(t *testing.T) { t.Fatal("daemon stop did not finish after health check unblocked") } + require.Eventually(t, func() bool { + _, ok := reg.Get("test-app", "serve") + return !ok + }, 500*time.Millisecond, 10*time.Millisecond, "daemon should unregister after health shutdown completes") + select { case err := <-healthErr: require.NoError(t, err)