From 73935333f641fe035715971fe2e1ef30c54135b7 Mon Sep 17 00:00:00 2001
From: Jana Radhakrishnan <mrjana@docker.com>
Date: Fri, 9 Oct 2015 19:31:44 -0700
Subject: [PATCH] Network ep count increment/decrement needs retry

Today we try to increment/decrement endpoint count
only once even if it is a key modified error. In case
of key modified error we should retry it to allow it to
succeed.

Signed-off-by: Jana Radhakrishnan <mrjana@docker.com>
---
 network.go | 38 ++++++++++++++++++++++++++++++--------
 store.go   |  3 +++
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/network.go b/network.go
index c4e80a4..ad1556a 100644
--- a/network.go
+++ b/network.go
@@ -316,20 +316,42 @@ func (n *network) EndpointCnt() uint64 {
 	return n.endpointCnt
 }
 
-func (n *network) IncEndpointCnt() error {
+func (n *network) atomicIncDecEpCnt(inc bool) error {
+retry:
 	n.Lock()
-	n.endpointCnt++
+	if inc {
+		n.endpointCnt++
+	} else {
+		n.endpointCnt--
+	}
 	n.Unlock()
 
-	return n.getController().updateToStore(n)
+	store := n.getController().getStore(n.DataScope())
+	if store == nil {
+		return fmt.Errorf("store not found for scope %s", n.DataScope())
+	}
+
+	if err := n.getController().updateToStore(n); err != nil {
+		if err == datastore.ErrKeyModified {
+			if err := store.GetObject(datastore.Key(n.Key()...), n); err != nil {
+				return fmt.Errorf("could not update the kvobject to latest when trying to atomic add endpoint count: %v", err)
+			}
+
+			goto retry
+		}
+
+		return err
+	}
+
+	return nil
+}
+
+func (n *network) IncEndpointCnt() error {
+	return n.atomicIncDecEpCnt(true)
 }
 
 func (n *network) DecEndpointCnt() error {
-	n.Lock()
-	n.endpointCnt--
-	n.Unlock()
-
-	return n.getController().updateToStore(n)
+	return n.atomicIncDecEpCnt(false)
 }
 
 // TODO : Can be made much more generic with the help of reflection (but has some golang limitations)
diff --git a/store.go b/store.go
index a45aa3b..c308423 100644
--- a/store.go
+++ b/store.go
@@ -155,6 +155,9 @@ func (c *controller) updateToStore(kvObject datastore.KVObject) error {
 	}
 
 	if err := cs.PutObjectAtomic(kvObject); err != nil {
+		if err == datastore.ErrKeyModified {
+			return err
+		}
 		return fmt.Errorf("failed to update store for object type %T: %v", kvObject, err)
 	}