Merge pull request #1539 from stevvooe/shutdown-session-on-update-sta…

…tus-failure

agent: bounce session after failed status update

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")

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
 }