From 353423b23ae159ba6ca39e03d2dba5c2d5208cb8 Mon Sep 17 00:00:00 2001
From: Stephen J Day <stephen.day@docker.com>
Date: Wed, 14 Sep 2016 16:57:52 -0700
Subject: [PATCH] agent: bounce session after failed status update

Previously, failed calls on an agent session to UpdateTaskStatus weren't
causing the agent to fall back into the session creation retry loop.
Without the retry loop, the agent ends up pounding the dispatcher with
status updates that will never drain.

This PR closes the session on a fatal error, ensuring that the agent
will fallback into a retry loop.

Signed-off-by: Stephen J Day <stephen.day@docker.com>
---
 agent/agent.go   |  3 ++-
 agent/session.go | 13 +++++++------
 2 files changed, 9 insertions(+), 7 deletions(-)

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
 }