@@ -243,7 +243,7 @@ RUN set -x \

 	&& rm -rf "$GOPATH"

 # Install containerd

-ENV CONTAINERD_COMMIT 4c21ad662f71af56c0e6b29c0afef72df441d1ff

+ENV CONTAINERD_COMMIT 2545227b0357eb55e369fa0072baef9ad91cdb69

 RUN set -x \

 	&& export GOPATH="$(mktemp -d)" \

 	&& git clone https://github.com/docker/containerd.git "$GOPATH/src/github.com/docker/containerd" \

@@ -186,7 +186,7 @@ RUN set -x \

 	&& rm -rf "$GOPATH"

 # Install containerd

-ENV CONTAINERD_COMMIT 4c21ad662f71af56c0e6b29c0afef72df441d1ff

+ENV CONTAINERD_COMMIT 2545227b0357eb55e369fa0072baef9ad91cdb69

 RUN set -x \

 	&& export GOPATH="$(mktemp -d)" \

 	&& git clone https://github.com/docker/containerd.git "$GOPATH/src/github.com/docker/containerd" \

@@ -184,7 +184,7 @@ RUN set -x \

 	&& rm -rf "$GOPATH"

 # Install containerd

-ENV CONTAINERD_COMMIT 4c21ad662f71af56c0e6b29c0afef72df441d1ff

+ENV CONTAINERD_COMMIT 2545227b0357eb55e369fa0072baef9ad91cdb69

 RUN set -x \

 	&& export GOPATH="$(mktemp -d)" \

 	&& git clone https://github.com/docker/containerd.git "$GOPATH/src/github.com/docker/containerd" \

@@ -204,7 +204,7 @@ RUN set -x \

 	&& rm -rf "$GOPATH"

 # Install containerd

-ENV CONTAINERD_COMMIT 4c21ad662f71af56c0e6b29c0afef72df441d1ff

+ENV CONTAINERD_COMMIT 2545227b0357eb55e369fa0072baef9ad91cdb69

 RUN set -x \

 	&& export GOPATH="$(mktemp -d)" \

 	&& git clone https://github.com/docker/containerd.git "$GOPATH/src/github.com/docker/containerd" \

@@ -196,7 +196,7 @@ RUN set -x \

 	&& rm -rf "$GOPATH"

 # Install containerd

-ENV CONTAINERD_COMMIT 4c21ad662f71af56c0e6b29c0afef72df441d1ff

+ENV CONTAINERD_COMMIT 2545227b0357eb55e369fa0072baef9ad91cdb69

 RUN set -x \

 	&& export GOPATH="$(mktemp -d)" \

 	&& git clone https://github.com/docker/containerd.git "$GOPATH/src/github.com/docker/containerd" \

@@ -68,7 +68,7 @@ RUN set -x \

 	&& rm -rf "$GOPATH"

 # Install containerd

-ENV CONTAINERD_COMMIT 4c21ad662f71af56c0e6b29c0afef72df441d1ff

+ENV CONTAINERD_COMMIT 2545227b0357eb55e369fa0072baef9ad91cdb69

 RUN set -x \

 	&& export GOPATH="$(mktemp -d)" \

 	&& git clone https://github.com/docker/containerd.git "$GOPATH/src/github.com/docker/containerd" \

@@ -141,10 +141,10 @@ clone git google.golang.org/cloud dae7e3d993bc3812a2185af60552bb6b847e52a0 https

 clone git github.com/docker/docker-credential-helpers v0.3.0

 # containerd

-clone git github.com/docker/containerd 4c21ad662f71af56c0e6b29c0afef72df441d1ff

+clone git github.com/docker/containerd 2545227b0357eb55e369fa0072baef9ad91cdb69

 # cluster

-clone git github.com/docker/swarmkit 27fbaef4ceed648bb575969ccc9083a6e104a719

+clone git github.com/docker/swarmkit 191acc1bbdb13d8ea3b8059dda14a12f8c3903f2

 clone git github.com/golang/mock bd3c8e81be01eef76d4b503f5e687d2d1354d2d9

 clone git github.com/gogo/protobuf v0.3

 clone git github.com/cloudflare/cfssl 7fb22c8cba7ecaf98e4082d22d65800cf45e042a

@@ -75,7 +75,9 @@ var _ = math.Inf

 // This is a compile-time assertion to ensure that this generated file

 // is compatible with the proto package it is being compiled against.

-const _ = proto.ProtoPackageIsVersion1

+// A compilation error at this line likely means your copy of the

+// proto package needs to be updated.

+const _ = proto.ProtoPackageIsVersion2 // please upgrade the proto package

 type GetServerVersionRequest struct {

 }

@@ -223,7 +225,7 @@ func (*Rlimit) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{9} }

 type User struct {

 	Uid            uint32   `protobuf:"varint,1,opt,name=uid" json:"uid,omitempty"`

 	Gid            uint32   `protobuf:"varint,2,opt,name=gid" json:"gid,omitempty"`

-	AdditionalGids []uint32 `protobuf:"varint,3,rep,name=additionalGids" json:"additionalGids,omitempty"`

+	AdditionalGids []uint32 `protobuf:"varint,3,rep,packed,name=additionalGids" json:"additionalGids,omitempty"`

 }

 func (m *User) Reset()                    { *m = User{} }

@@ -385,7 +387,7 @@ type Container struct {

 	Processes  []*Process `protobuf:"bytes,3,rep,name=processes" json:"processes,omitempty"`

 	Status     string     `protobuf:"bytes,4,opt,name=status" json:"status,omitempty"`

 	Labels     []string   `protobuf:"bytes,5,rep,name=labels" json:"labels,omitempty"`

-	Pids       []uint32   `protobuf:"varint,6,rep,name=pids" json:"pids,omitempty"`

+	Pids       []uint32   `protobuf:"varint,6,rep,packed,name=pids" json:"pids,omitempty"`

 	Runtime    string     `protobuf:"bytes,7,opt,name=runtime" json:"runtime,omitempty"`

 }

@@ -628,7 +630,7 @@ func (*NetworkStats) Descriptor() ([]byte, []int) { return fileDescriptor0, []in

 type CpuUsage struct {

 	TotalUsage        uint64   `protobuf:"varint,1,opt,name=total_usage,json=totalUsage" json:"total_usage,omitempty"`

-	PercpuUsage       []uint64 `protobuf:"varint,2,rep,name=percpu_usage,json=percpuUsage" json:"percpu_usage,omitempty"`

+	PercpuUsage       []uint64 `protobuf:"varint,2,rep,packed,name=percpu_usage,json=percpuUsage" json:"percpu_usage,omitempty"`

 	UsageInKernelmode uint64   `protobuf:"varint,3,opt,name=usage_in_kernelmode,json=usageInKernelmode" json:"usage_in_kernelmode,omitempty"`

 	UsageInUsermode   uint64   `protobuf:"varint,4,opt,name=usage_in_usermode,json=usageInUsermode" json:"usage_in_usermode,omitempty"`

 }

@@ -978,7 +980,7 @@ var _ grpc.ClientConn

 // This is a compile-time assertion to ensure that this generated file

 // is compatible with the grpc package it is being compiled against.

-const _ = grpc.SupportPackageIsVersion2

+const _ = grpc.SupportPackageIsVersion3

 // Client API for API service

@@ -1432,8 +1434,11 @@ var _API_serviceDesc = grpc.ServiceDesc{

 			ServerStreams: true,

 		},

 	},

+	Metadata: fileDescriptor0,

 }

+func init() { proto.RegisterFile("api.proto", fileDescriptor0) }

+

 var fileDescriptor0 = []byte{

 	// 2604 bytes of a gzipped FileDescriptorProto

 	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x02, 0xff, 0xec, 0x59, 0x4b, 0x6f, 0x1c, 0x5b,

deleted file mode 100644

@@ -1,34 +0,0 @@

-package reference

-

-import (

-	distreference "github.com/docker/distribution/reference"

-)

-

-// Parse parses the given references and returns the repository and

-// tag (if present) from it. If there is an error during parsing, it will

-// return an error.

-func Parse(ref string) (string, string, error) {

-	distributionRef, err := distreference.ParseNamed(ref)

-	if err != nil {

-		return "", "", err

-	}

-

-	tag := GetTagFromNamedRef(distributionRef)

-	return distributionRef.Name(), tag, nil

-}

-

-// GetTagFromNamedRef returns a tag from the specified reference.

-// This function is necessary as long as the docker "server" api makes the distinction between repository

-// and tags.

-func GetTagFromNamedRef(ref distreference.Named) string {

-	var tag string

-	switch x := ref.(type) {

-	case distreference.Digested:

-		tag = x.Digest().String()

-	case distreference.NamedTagged:

-		tag = x.Tag()

-	default:

-		tag = "latest"

-	}

-	return tag

-}

@@ -15,6 +15,7 @@ import (

 const (

 	initialSessionFailureBackoff = 100 * time.Millisecond

 	maxSessionFailureBackoff     = 8 * time.Second

+	nodeUpdatePeriod             = 20 * time.Second

 )

 // Agent implements the primary node functionality for a member of a swarm

@@ -134,9 +135,18 @@ func (a *Agent) run(ctx context.Context) {

 	log.G(ctx).Debugf("(*Agent).run")

 	defer log.G(ctx).Debugf("(*Agent).run exited")

+	// get the node description

+	nodeDescription, err := a.nodeDescriptionWithHostname(ctx)

+	if err != nil {

+		log.G(ctx).WithError(err).WithField("agent", a.config.Executor).Errorf("agent: node description unavailable")

+	}

+	// nodeUpdateTicker is used to periodically check for updates to node description

+	nodeUpdateTicker := time.NewTicker(nodeUpdatePeriod)

+	defer nodeUpdateTicker.Stop()

+

 	var (

 		backoff    time.Duration

-		session    = newSession(ctx, a, backoff, "") // start the initial session

+		session    = newSession(ctx, a, backoff, "", nodeDescription) // start the initial session

 		registered = session.registered

 		ready      = a.ready // first session ready

 		sessionq   chan sessionOperation

@@ -158,9 +168,16 @@ func (a *Agent) run(ctx context.Context) {

 		select {

 		case operation := <-sessionq:

 			operation.response <- operation.fn(session)

-		case msg := <-session.tasks:

-			if err := a.worker.Assign(ctx, msg.Tasks); err != nil {

-				log.G(ctx).WithError(err).Error("task assignment failed")

+		case msg := <-session.assignments:

+			switch msg.Type {

+			case api.AssignmentsMessage_COMPLETE:

+				if err := a.worker.AssignTasks(ctx, msg.UpdateTasks); err != nil {

+					log.G(ctx).WithError(err).Error("failed to synchronize worker assignments")

+				}

+			case api.AssignmentsMessage_INCREMENTAL:

+				if err := a.worker.UpdateTasks(ctx, msg.UpdateTasks, msg.RemoveTasks); err != nil {

+					log.G(ctx).WithError(err).Error("failed to update worker assignments")

+				}

 			}

 		case msg := <-session.messages:

 			if err := a.handleSessionMessage(ctx, msg); err != nil {

@@ -197,10 +214,42 @@ func (a *Agent) run(ctx context.Context) {

 			log.G(ctx).Debugf("agent: rebuild session")

 			// select a session registration delay from backoff range.

-			delay := time.Duration(rand.Int63n(int64(backoff)))

-			session = newSession(ctx, a, delay, session.sessionID)

+			delay := time.Duration(0)

+			if backoff > 0 {

+				delay = time.Duration(rand.Int63n(int64(backoff)))

+			}

+			session = newSession(ctx, a, delay, session.sessionID, nodeDescription)

 			registered = session.registered

 			sessionq = a.sessionq

+		case <-nodeUpdateTicker.C:

+			// skip this case if the registration isn't finished

+			if registered != nil {

+				continue

+			}

+			// get the current node description

+			newNodeDescription, err := a.nodeDescriptionWithHostname(ctx)

+			if err != nil {

+				log.G(ctx).WithError(err).WithField("agent", a.config.Executor).Errorf("agent: updated node description unavailable")

+			}

+

+			// if newNodeDescription is nil, it will cause a panic when

+			// trying to create a session. Typically this can happen

+			// if the engine goes down

+			if newNodeDescription == nil {

+				continue

+			}

+

+			// if the node description has changed, update it to the new one

+			// and close the session. The old session will be stopped and a

+			// new one will be created with the updated description

+			if !reflect.DeepEqual(nodeDescription, newNodeDescription) {

+				nodeDescription = newNodeDescription

+				// close the session

+				log.G(ctx).Info("agent: found node update")

+				if err := session.close(); err != nil {

+					log.G(ctx).WithError(err).Error("agent: closing session for node update failed")

+				}

+			}

 		case <-a.stopped:

 			// TODO(stevvooe): Wait on shutdown and cleanup. May need to pump

 			// this loop a few times.

@@ -315,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")

@@ -337,6 +387,17 @@ func (a *Agent) UpdateTaskStatus(ctx context.Context, taskID string, status *api

 	}

 }

+// nodeDescriptionWithHostname retrieves node description, and overrides hostname if available

+func (a *Agent) nodeDescriptionWithHostname(ctx context.Context) (*api.NodeDescription, error) {

+	desc, err := a.config.Executor.Describe(ctx)

+

+	// Override hostname

+	if a.config.Hostname != "" && desc != nil {

+		desc.Hostname = a.config.Hostname

+	}

+	return desc, err

+}

+

 // nodesEqual returns true if the node states are functionaly equal, ignoring status,

 // version and other superfluous fields.

 //

@@ -29,7 +29,7 @@ type Config struct {

 	NotifyRoleChange chan<- api.NodeRole

 	// Credentials is credentials for grpc connection to manager.

-	Credentials credentials.TransportAuthenticator

+	Credentials credentials.TransportCredentials

 }

 func (c *Config) validate() error {

@@ -147,7 +147,7 @@ func Do(ctx context.Context, task *api.Task, ctlr Controller) (*api.TaskStatus,

 		if cs, ok := err.(ContainerStatuser); ok {

 			var err error

 			containerStatus, err = cs.ContainerStatus(ctx)

-			if err != nil {

+			if err != nil && !contextDoneError(err) {

 				log.G(ctx).WithError(err).Error("error resolving container status on fatal")

 			}

 		}

@@ -207,7 +207,7 @@ func Do(ctx context.Context, task *api.Task, ctlr Controller) (*api.TaskStatus,

 			var err error

 			containerStatus, err = cctlr.ContainerStatus(ctx)

-			if err != nil {

+			if err != nil && !contextDoneError(err) {

 				log.G(ctx).WithError(err).Error("container status unavailable")

 			}

@@ -297,3 +297,8 @@ func logStateChange(ctx context.Context, desired, previous, next api.TaskState)

 		log.G(ctx).WithFields(fields).Debug("state changed")

 	}

 }

+

+func contextDoneError(err error) bool {

+	cause := errors.Cause(err)

+	return cause == context.Canceled || cause == context.DeadlineExceeded

+}

@@ -120,7 +120,7 @@ func NewNode(c *NodeConfig) (*Node, error) {

 	n := &Node{

 		remotes:              newPersistentRemotes(stateFile, p...),

-		role:                 ca.AgentRole,

+		role:                 ca.WorkerRole,

 		config:               c,

 		started:              make(chan struct{}),

 		stopped:              make(chan struct{}),

@@ -194,7 +194,9 @@ func (n *Node) run(ctx context.Context) (err error) {

 		select {

 		case <-ctx.Done():

 		case resp := <-issueResponseChan:

-			logrus.Debugf("Requesting certificate for NodeID: %v", resp.NodeID)

+			log.G(log.WithModule(ctx, "tls")).WithFields(logrus.Fields{

+				"node.id": resp.NodeID,

+			}).Debugf("requesting certificate")

 			n.Lock()

 			n.nodeID = resp.NodeID

 			n.nodeMembership = resp.NodeMembership

@@ -233,7 +235,7 @@ func (n *Node) run(ctx context.Context) (err error) {

 			case apirole := <-n.roleChangeReq:

 				n.Lock()

 				lastRole := n.role

-				role := ca.AgentRole

+				role := ca.WorkerRole

 				if apirole == api.NodeRoleManager {

 					role = ca.ManagerRole

 				}

@@ -242,7 +244,7 @@ func (n *Node) run(ctx context.Context) (err error) {

 					continue

 				}

 				// switch role to agent immediately to shutdown manager early

-				if role == ca.AgentRole {

+				if role == ca.WorkerRole {

 					n.role = role

 					n.roleCond.Broadcast()

 				}

@@ -343,7 +345,7 @@ func (n *Node) Err(ctx context.Context) error {

 	}

 }

-func (n *Node) runAgent(ctx context.Context, db *bolt.DB, creds credentials.TransportAuthenticator, ready chan<- struct{}) error {

+func (n *Node) runAgent(ctx context.Context, db *bolt.DB, creds credentials.TransportCredentials, ready chan<- struct{}) error {

 	select {

 	case <-ctx.Done():

 	case <-n.remotes.WaitSelect(ctx):

@@ -588,7 +590,7 @@ func (n *Node) runManager(ctx context.Context, securityConfig *ca.SecurityConfig

 			return err

 		}

-		remoteAddr, _ := n.remotes.Select(n.nodeID)

+		remoteAddr, _ := n.remotes.Select(n.NodeID())

 		m, err := manager.New(&manager.Config{

 			ForceNewCluster: n.config.ForceNewCluster,

 			ProtoAddr: map[string]string{

@@ -607,8 +609,9 @@ func (n *Node) runManager(ctx context.Context, securityConfig *ca.SecurityConfig

 			return err

 		}

 		done := make(chan struct{})

+		var runErr error

 		go func() {

-			m.Run(context.Background()) // todo: store error

+			runErr = m.Run(context.Background())

 			close(done)

 		}()

@@ -624,14 +627,31 @@ func (n *Node) runManager(ctx context.Context, securityConfig *ca.SecurityConfig

 			go func(ready chan struct{}) {

 				select {

 				case <-ready:

-					n.remotes.Observe(api.Peer{NodeID: n.nodeID, Addr: n.config.ListenRemoteAPI}, remotes.DefaultObservationWeight)

+					n.remotes.Observe(api.Peer{NodeID: n.NodeID(), Addr: n.config.ListenRemoteAPI}, remotes.DefaultObservationWeight)

 				case <-connCtx.Done():

 				}

 			}(ready)

 			ready = nil

 		}

-		err = n.waitRole(ctx, ca.AgentRole)

+		roleChanged := make(chan error)

+		waitCtx, waitCancel := context.WithCancel(ctx)

+		go func() {

+			err := n.waitRole(waitCtx, ca.WorkerRole)

+			roleChanged <- err

+		}()

+

+		select {

+		case <-done:

+			// Fail out if m.Run() returns error, otherwise wait for

+			// role change.

+			if runErr != nil {

+				err = runErr

+			} else {

+				err = <-roleChanged

+			}

+		case err = <-roleChanged:

+		}

 		n.Lock()

 		n.manager = nil

@@ -646,6 +666,7 @@ func (n *Node) runManager(ctx context.Context, securityConfig *ca.SecurityConfig

 		}

 		connCancel()

 		n.setControlSocket(nil)

+		waitCancel()

 		if err != nil {

 			return err

@@ -672,17 +693,18 @@ func newPersistentRemotes(f string, peers ...api.Peer) *persistentRemotes {

 func (s *persistentRemotes) Observe(peer api.Peer, weight int) {

 	s.Lock()

+	defer s.Unlock()

 	s.Remotes.Observe(peer, weight)

 	s.c.Broadcast()

 	if err := s.save(); err != nil {

 		logrus.Errorf("error writing cluster state file: %v", err)

-		s.Unlock()

 		return

 	}

-	s.Unlock()

 	return

 }

 func (s *persistentRemotes) Remove(peers ...api.Peer) {

+	s.Lock()

+	defer s.Unlock()

 	s.Remotes.Remove(peers...)

 	if err := s.save(); err != nil {

 		logrus.Errorf("error writing cluster state file: %v", err)

@@ -2,8 +2,10 @@ package agent

 import (

 	"errors"

+	"sync"

 	"time"

+	"github.com/Sirupsen/logrus"

 	"github.com/docker/swarmkit/api"

 	"github.com/docker/swarmkit/log"

 	"github.com/docker/swarmkit/protobuf/ptypes"

@@ -31,26 +33,27 @@ type session struct {

 	conn *grpc.ClientConn

 	addr string

-	agent     *Agent

-	sessionID string

-	session   api.Dispatcher_SessionClient

-	errs      chan error

-	messages  chan *api.SessionMessage

-	tasks     chan *api.TasksMessage

+	agent       *Agent

+	sessionID   string

+	session     api.Dispatcher_SessionClient

+	errs        chan error

+	messages    chan *api.SessionMessage

+	assignments chan *api.AssignmentsMessage

 	registered chan struct{} // closed registration

 	closed     chan struct{}

+	closeOnce  sync.Once

 }

-func newSession(ctx context.Context, agent *Agent, delay time.Duration, sessionID string) *session {

+func newSession(ctx context.Context, agent *Agent, delay time.Duration, sessionID string, description *api.NodeDescription) *session {

 	s := &session{

-		agent:      agent,

-		sessionID:  sessionID,

-		errs:       make(chan error, 1),

-		messages:   make(chan *api.SessionMessage),

-		tasks:      make(chan *api.TasksMessage),

-		registered: make(chan struct{}),

-		closed:     make(chan struct{}),

+		agent:       agent,

+		sessionID:   sessionID,

+		errs:        make(chan error, 1),

+		messages:    make(chan *api.SessionMessage),

+		assignments: make(chan *api.AssignmentsMessage),

+		registered:  make(chan struct{}),

+		closed:      make(chan struct{}),

 	}

 	peer, err := agent.config.Managers.Select()

 	if err != nil {

@@ -68,14 +71,14 @@ func newSession(ctx context.Context, agent *Agent, delay time.Duration, sessionI

 	s.addr = peer.Addr

 	s.conn = cc

-	go s.run(ctx, delay)

+	go s.run(ctx, delay, description)

 	return s

 }

-func (s *session) run(ctx context.Context, delay time.Duration) {

+func (s *session) run(ctx context.Context, delay time.Duration, description *api.NodeDescription) {

 	time.Sleep(delay) // delay before registering.

-	if err := s.start(ctx); err != nil {

+	if err := s.start(ctx, description); err != nil {

 		select {

 		case s.errs <- err:

 		case <-s.closed:

@@ -94,24 +97,14 @@ func (s *session) run(ctx context.Context, delay time.Duration) {

 }

 // start begins the session and returns the first SessionMessage.

-func (s *session) start(ctx context.Context) error {

+func (s *session) start(ctx context.Context, description *api.NodeDescription) error {

 	log.G(ctx).Debugf("(*session).start")

-	description, err := s.agent.config.Executor.Describe(ctx)

-	if err != nil {

-		log.G(ctx).WithError(err).WithField("executor", s.agent.config.Executor).

-			Errorf("node description unavailable")

-		return err

-	}

-	// Override hostname

-	if s.agent.config.Hostname != "" {

-		description.Hostname = s.agent.config.Hostname

-	}

-

 	errChan := make(chan error, 1)

 	var (

 		msg    *api.SessionMessage

 		stream api.Dispatcher_SessionClient

+		err    error

 	)

 	// Note: we don't defer cancellation of this context, because the

 	// streaming RPC is used after this function returned. We only cancel

@@ -215,22 +208,68 @@ func (s *session) handleSessionMessage(ctx context.Context, msg *api.SessionMess

 }

 func (s *session) watch(ctx context.Context) error {

-	log.G(ctx).Debugf("(*session).watch")

-	client := api.NewDispatcherClient(s.conn)

-	watch, err := client.Tasks(ctx, &api.TasksRequest{

-		SessionID: s.sessionID})

-	if err != nil {

-		return err

-	}

+	log := log.G(ctx).WithFields(logrus.Fields{"method": "(*session).watch"})

+	log.Debugf("")

+	var (

+		resp            *api.AssignmentsMessage

+		assignmentWatch api.Dispatcher_AssignmentsClient

+		tasksWatch      api.Dispatcher_TasksClient

+		streamReference string

+		tasksFallback   bool

+		err             error

+	)

+	client := api.NewDispatcherClient(s.conn)

 	for {

-		resp, err := watch.Recv()

-		if err != nil {

-			return err

+		// If this is the first time we're running the loop, or there was a reference mismatch

+		// attempt to get the assignmentWatch

+		if assignmentWatch == nil && !tasksFallback {

+			assignmentWatch, err = client.Assignments(ctx, &api.AssignmentsRequest{SessionID: s.sessionID})

+			if err != nil {

+				return err

+			}

+		}

+		// We have an assignmentWatch, let's try to receive an AssignmentMessage

+		if assignmentWatch != nil {

+			// If we get a code = 12 desc = unknown method Assignments, try to use tasks

+			resp, err = assignmentWatch.Recv()

+			if err != nil {

+				if grpc.Code(err) != codes.Unimplemented {

+					return err

+				}

+				tasksFallback = true

+				assignmentWatch = nil

+				log.WithError(err).Infof("falling back to Tasks")

+			}

+		}

+

+		// This code is here for backwards compatibility (so that newer clients can use the

+		// older method Tasks)

+		if tasksWatch == nil && tasksFallback {

+			tasksWatch, err = client.Tasks(ctx, &api.TasksRequest{SessionID: s.sessionID})

+			if err != nil {

+				return err

+			}

+		}

+		if tasksWatch != nil {

+			var taskResp *api.TasksMessage

+			taskResp, err = tasksWatch.Recv()

+			if err != nil {

+				return err

+			}

+			resp = &api.AssignmentsMessage{Type: api.AssignmentsMessage_COMPLETE, UpdateTasks: taskResp.Tasks}

+		}

+

+		// If there seems to be a gap in the stream, let's break out of the inner for and

+		// re-sync (by calling Assignments again).

+		if streamReference != "" && streamReference != resp.AppliesTo {

+			assignmentWatch = nil

+		} else {

+			streamReference = resp.ResultsIn

 		}

 		select {

-		case s.tasks <- resp:

+		case s.assignments <- resp:

 		case <-s.closed:

 			return errSessionClosed

 		case <-ctx.Done():

@@ -241,7 +280,6 @@ func (s *session) watch(ctx context.Context) error {

 // sendTaskStatus uses the current session to send the status of a single task.

 func (s *session) sendTaskStatus(ctx context.Context, taskID string, status *api.TaskStatus) error {

-

 	client := api.NewDispatcherClient(s.conn)

 	if _, err := client.UpdateTaskStatus(ctx, &api.UpdateTaskStatusRequest{

 		SessionID: s.sessionID,

@@ -302,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

 }

@@ -17,9 +17,13 @@ type Worker interface {

 	// Init prepares the worker for task assignment.

 	Init(ctx context.Context) error

-	// Assign the set of tasks to the worker. Tasks outside of this set will be

-	// removed.

-	Assign(ctx context.Context, tasks []*api.Task) error

+	// AssignTasks assigns a complete set of tasks to a worker. Any task not included in

+	// this set will be removed.

+	AssignTasks(ctx context.Context, tasks []*api.Task) error

+

+	// UpdateTasks updates an incremental set of tasks to the worker. Any task not included

+	// either in added or removed will remain untouched.

+	UpdateTasks(ctx context.Context, added []*api.Task, removed []string) error

 	// Listen to updates about tasks controlled by the worker. When first

 	// called, the reporter will receive all updates for all tasks controlled

@@ -86,14 +90,37 @@ func (w *worker) Init(ctx context.Context) error {

 	})

 }

-// Assign the set of tasks to the worker. Any tasks not previously known will

+// AssignTasks assigns  the set of tasks to the worker. Any tasks not previously known will

 // be started. Any tasks that are in the task set and already running will be

 // updated, if possible. Any tasks currently running on the

 // worker outside the task set will be terminated.

-func (w *worker) Assign(ctx context.Context, tasks []*api.Task) error {

+func (w *worker) AssignTasks(ctx context.Context, tasks []*api.Task) error {

 	w.mu.Lock()

 	defer w.mu.Unlock()

+	log.G(ctx).WithFields(logrus.Fields{

+		"len(tasks)": len(tasks),

+	}).Debug("(*worker).AssignTasks")

+

+	return reconcileTaskState(ctx, w, tasks, nil, true)

+}

+

+// UpdateTasks the set of tasks to the worker.

+// Tasks in the added set will be added to the worker, and tasks in the removed set

+// will be removed from the worker

+func (w *worker) UpdateTasks(ctx context.Context, added []*api.Task, removed []string) error {

+	w.mu.Lock()

+	defer w.mu.Unlock()

+

+	log.G(ctx).WithFields(logrus.Fields{

+		"len(added)":   len(added),

+		"len(removed)": len(removed),

+	}).Debug("(*worker).UpdateTasks")

+

+	return reconcileTaskState(ctx, w, added, removed, false)

+}

+

+func reconcileTaskState(ctx context.Context, w *worker, added []*api.Task, removed []string, fullSnapshot bool) error {

 	tx, err := w.db.Begin(true)

 	if err != nil {

 		log.G(ctx).WithError(err).Error("failed starting transaction against task database")

@@ -101,10 +128,9 @@ func (w *worker) Assign(ctx context.Context, tasks []*api.Task) error {

 	}

 	defer tx.Rollback()

-	log.G(ctx).WithField("len(tasks)", len(tasks)).Debug("(*worker).Assign")

 	assigned := map[string]struct{}{}

-	for _, task := range tasks {

+	for _, task := range added {

 		log.G(ctx).WithFields(

 			logrus.Fields{

 				"task.id":           task.ID,

@@ -135,35 +161,59 @@ func (w *worker) Assign(ctx context.Context, tasks []*api.Task) error {

 					return err

 				}

 			} else {

-				task.Status = *status // overwrite the stale manager status with ours.

+				task.Status = *status

 			}

-

 			w.startTask(ctx, tx, task)

 		}

 		assigned[task.ID] = struct{}{}

 	}

-	for id, tm := range w.taskManagers {

-		if _, ok := assigned[id]; ok {

-			continue

+	closeManager := func(tm *taskManager) {

+		// when a task is no longer assigned, we shutdown the task manager for

+		// it and leave cleanup to the sweeper.

+		if err := tm.Close(); err != nil {

+			log.G(ctx).WithError(err).Error("error closing task manager")

 		}

+	}

-		ctx := log.WithLogger(ctx, log.G(ctx).WithField("task.id", id))

-		if err := SetTaskAssignment(tx, id, false); err != nil {

+	removeTaskAssignment := func(taskID string) error {

+		ctx := log.WithLogger(ctx, log.G(ctx).WithField("task.id", taskID))

+		if err := SetTaskAssignment(tx, taskID, false); err != nil {

 			log.G(ctx).WithError(err).Error("error setting task assignment in database")

-			continue

 		}

+		return err

+	}

+

+	// If this was a complete set of assignments, we're going to remove all the remaining

+	// tasks.

+	if fullSnapshot {

+		for id, tm := range w.taskManagers {

+			if _, ok := assigned[id]; ok {

+				continue

+			}

-		delete(w.taskManagers, id)

+			err := removeTaskAssignment(id)

+			if err == nil {

+				delete(w.taskManagers, id)

+				go closeManager(tm)

+			}

+		}

+	} else {

+		// If this was an incremental set of assignments, we're going to remove only the tasks

+		// in the removed set

+		for _, taskID := range removed {

+			err := removeTaskAssignment(taskID)

+			if err != nil {

+				continue

+			}

-		go func(tm *taskManager) {

-			// when a task is no longer assigned, we shutdown the task manager for

-			// it and leave cleanup to the sweeper.

-			if err := tm.Close(); err != nil {

-				log.G(ctx).WithError(err).Error("error closing task manager")

+			tm, ok := w.taskManagers[taskID]

+			if ok {

+				delete(w.taskManagers, taskID)

+				go closeManager(tm)

 			}

-		}(tm)

+		}

 	}

 	return tx.Commit()

@@ -21,10 +21,11 @@ import (

 	grpc "google.golang.org/grpc"

 )

-import raftpicker "github.com/docker/swarmkit/manager/raftpicker"

+import raftselector "github.com/docker/swarmkit/manager/raftselector"

 import codes "google.golang.org/grpc/codes"

 import metadata "google.golang.org/grpc/metadata"

 import transport "google.golang.org/grpc/transport"

+import time "time"

 import io "io"

@@ -285,11 +286,12 @@ func valueToGoStringCa(v interface{}, typ string) string {

 	pv := reflect.Indirect(rv).Interface()

 	return fmt.Sprintf("func(v %v) *%v { return &v } ( %#v )", typ, typ, pv)

 }

-func extensionToGoStringCa(e map[int32]github_com_gogo_protobuf_proto.Extension) string {

+func extensionToGoStringCa(m github_com_gogo_protobuf_proto.Message) string {

+	e := github_com_gogo_protobuf_proto.GetUnsafeExtensionsMap(m)

 	if e == nil {

 		return "nil"

 	}

-	s := "map[int32]proto.Extension{"

+	s := "proto.NewUnsafeXXX_InternalExtensions(map[int32]proto.Extension{"

 	keys := make([]int, 0, len(e))

 	for k := range e {

 		keys = append(keys, int(k))

@@ -299,7 +301,7 @@ func extensionToGoStringCa(e map[int32]github_com_gogo_protobuf_proto.Extension)

 	for _, k := range keys {

 		ss = append(ss, strconv.Itoa(k)+": "+e[int32(k)].GoString())

 	}

-	s += strings.Join(ss, ",") + "}"

+	s += strings.Join(ss, ",") + "})"

 	return s

 }

@@ -309,7 +311,7 @@ var _ grpc.ClientConn

 // This is a compile-time assertion to ensure that this generated file

 // is compatible with the grpc package it is being compiled against.

-const _ = grpc.SupportPackageIsVersion2

+const _ = grpc.SupportPackageIsVersion3

 // Client API for CA service

@@ -371,7 +373,8 @@ var _CA_serviceDesc = grpc.ServiceDesc{

 			Handler:    _CA_GetRootCACertificate_Handler,

 		},

 	},

-	Streams: []grpc.StreamDesc{},

+	Streams:  []grpc.StreamDesc{},

+	Metadata: fileDescriptorCa,

 }

 // Client API for NodeCA service

@@ -467,7 +470,8 @@ var _NodeCA_serviceDesc = grpc.ServiceDesc{