graceful shutdown fix

This commit is contained in:
Grigoriy Mikhalkin 2022-06-30 23:35:22 +02:00
parent c6eb7be7fb
commit 889eff265f
2 changed files with 29 additions and 15 deletions

21
poll.go
View file

@ -290,6 +290,9 @@ func (h *Headscale) PollNetMapStream(
keepAliveChan chan []byte,
updateChan chan struct{},
) {
h.wg.Add(1)
defer h.wg.Done()
ctx := context.WithValue(req.Context(), machineNameContextKey, machine.Hostname)
ctx, cancel := context.WithCancel(ctx)
@ -353,9 +356,9 @@ func (h *Headscale) PollNetMapStream(
Str("channel", "pollData").
Int("bytes", len(data)).
Msg("Data from pollData channel written successfully")
// TODO(kradalby): Abstract away all the database calls, this can cause race conditions
// when an outdated machine object is kept alive, e.g. db is update from
// command line, but then overwritten.
// TODO(kradalby): Abstract away all the database calls, this can cause race conditions
// when an outdated machine object is kept alive, e.g. db is update from
// command line, but then overwritten.
err = h.UpdateMachineFromDatabase(machine)
if err != nil {
log.Error().
@ -431,9 +434,9 @@ func (h *Headscale) PollNetMapStream(
Str("channel", "keepAlive").
Int("bytes", len(data)).
Msg("Keep alive sent successfully")
// TODO(kradalby): Abstract away all the database calls, this can cause race conditions
// when an outdated machine object is kept alive, e.g. db is update from
// command line, but then overwritten.
// TODO(kradalby): Abstract away all the database calls, this can cause race conditions
// when an outdated machine object is kept alive, e.g. db is update from
// command line, but then overwritten.
err = h.UpdateMachineFromDatabase(machine)
if err != nil {
log.Error().
@ -588,9 +591,9 @@ func (h *Headscale) PollNetMapStream(
Str("handler", "PollNetMapStream").
Str("machine", machine.Hostname).
Msg("The client has closed the connection")
// TODO: Abstract away all the database calls, this can cause race conditions
// when an outdated machine object is kept alive, e.g. db is update from
// command line, but then overwritten.
// TODO: Abstract away all the database calls, this can cause race conditions
// when an outdated machine object is kept alive, e.g. db is update from
// command line, but then overwritten.
err := h.UpdateMachineFromDatabase(machine)
if err != nil {
log.Error().