diff --git a/agent/agent.go b/agent/agent.go index 5972f8cb70..8f686a0a07 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -364,7 +364,8 @@ func (a *Agent) UpdateTaskStatus(ctx context.Context, taskID string, status *api if err == errTaskUnknown { err = nil // dispatcher no longer cares about this task. } else { - log.G(ctx).WithError(err).Error("sending task status update failed") + log.G(ctx).WithError(err).Error("closing session after fatal error") + session.close() } } else { log.G(ctx).Debug("task status reported") diff --git a/agent/session.go b/agent/session.go index b4afc0d6f0..fc1a4582ce 100644 --- a/agent/session.go +++ b/agent/session.go @@ -2,6 +2,7 @@ package agent import ( "errors" + "sync" "time" "github.com/Sirupsen/logrus" @@ -41,6 +42,7 @@ type session struct { registered chan struct{} // closed registration closed chan struct{} + closeOnce sync.Once } func newSession(ctx context.Context, agent *Agent, delay time.Duration, sessionID string, description *api.NodeDescription) *session { @@ -338,15 +340,14 @@ func (s *session) sendTaskStatuses(ctx context.Context, updates ...*api.UpdateTa } func (s *session) close() error { - select { - case <-s.closed: - return errSessionClosed - default: + s.closeOnce.Do(func() { if s.conn != nil { s.agent.config.Managers.ObserveIfExists(api.Peer{Addr: s.addr}, -remotes.DefaultObservationWeight) s.conn.Close() } + close(s.closed) - return nil - } + }) + + return nil }