From 6a92628ddbb021e9661d11f27fa8d736f3cfe620 Mon Sep 17 00:00:00 2001
From: Alexandre Beslic <abronan@docker.com>
Date: Wed, 27 May 2015 08:02:54 -0700
Subject: [PATCH] Add ErrKeyModified error to etcd and zookeeper

Signed-off-by: Alexandre Beslic <abronan@docker.com>
---
 pkg/store/consul.go    |  4 ++--
 pkg/store/etcd.go      | 14 ++++++++++++--
 pkg/store/zookeeper.go | 10 ++++++++--
 3 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/pkg/store/consul.go b/pkg/store/consul.go
index 523ce1de57..8f2c07e495 100644
--- a/pkg/store/consul.go
+++ b/pkg/store/consul.go
@@ -343,7 +343,7 @@ func (l *consulLock) Unlock() error {
 // AtomicPut put a value at "key" if the key has not been
 // modified in the meantime, throws an error if this is the case
 func (s *Consul) AtomicPut(key string, value []byte, previous *KVPair, options *WriteOptions) (bool, *KVPair, error) {
-	if previous != nil {
+	if previous == nil {
 		return false, nil, ErrPreviousNotSpecified
 	}
 
@@ -372,7 +372,7 @@ func (s *Consul) AtomicPut(key string, value []byte, previous *KVPair, options *
 // AtomicDelete deletes a value at "key" if the key has not
 // been modified in the meantime, throws an error if this is the case
 func (s *Consul) AtomicDelete(key string, previous *KVPair) (bool, error) {
-	if previous != nil {
+	if previous == nil {
 		return false, ErrPreviousNotSpecified
 	}
 
diff --git a/pkg/store/etcd.go b/pkg/store/etcd.go
index 1f794a145b..8e4f6a3785 100644
--- a/pkg/store/etcd.go
+++ b/pkg/store/etcd.go
@@ -253,12 +253,17 @@ func (s *Etcd) WatchTree(prefix string, stopCh <-chan struct{}) (<-chan []*KVPai
 // AtomicPut put a value at "key" if the key has not been
 // modified in the meantime, throws an error if this is the case
 func (s *Etcd) AtomicPut(key string, value []byte, previous *KVPair, options *WriteOptions) (bool, *KVPair, error) {
-	if previous != nil {
+	if previous == nil {
 		return false, nil, ErrPreviousNotSpecified
 	}
 
 	meta, err := s.client.CompareAndSwap(normalize(key), string(value), 0, "", previous.LastIndex)
 	if err != nil {
+		if etcdError, ok := err.(*etcd.EtcdError); ok {
+			if etcdError.ErrorCode == 101 { // Compare failed
+				return false, nil, ErrKeyModified
+			}
+		}
 		return false, nil, err
 	}
 	return true, &KVPair{Key: key, Value: value, LastIndex: meta.Node.ModifiedIndex}, nil
@@ -267,12 +272,17 @@ func (s *Etcd) AtomicPut(key string, value []byte, previous *KVPair, options *Wr
 // AtomicDelete deletes a value at "key" if the key has not
 // been modified in the meantime, throws an error if this is the case
 func (s *Etcd) AtomicDelete(key string, previous *KVPair) (bool, error) {
-	if previous != nil {
+	if previous == nil {
 		return false, ErrPreviousNotSpecified
 	}
 
 	_, err := s.client.CompareAndDelete(normalize(key), "", previous.LastIndex)
 	if err != nil {
+		if etcdError, ok := err.(*etcd.EtcdError); ok {
+			if etcdError.ErrorCode == 101 { // Compare failed
+				return false, ErrKeyModified
+			}
+		}
 		return false, err
 	}
 	return true, nil
diff --git a/pkg/store/zookeeper.go b/pkg/store/zookeeper.go
index d5989ad9e1..d7d977ac8d 100644
--- a/pkg/store/zookeeper.go
+++ b/pkg/store/zookeeper.go
@@ -215,12 +215,15 @@ func (s *Zookeeper) DeleteTree(prefix string) error {
 // AtomicPut put a value at "key" if the key has not been
 // modified in the meantime, throws an error if this is the case
 func (s *Zookeeper) AtomicPut(key string, value []byte, previous *KVPair, _ *WriteOptions) (bool, *KVPair, error) {
-	if previous != nil {
+	if previous == nil {
 		return false, nil, ErrPreviousNotSpecified
 	}
 
 	meta, err := s.client.Set(normalize(key), value, int32(previous.LastIndex))
 	if err != nil {
+		if err == zk.ErrBadVersion {
+			return false, nil, ErrKeyModified
+		}
 		return false, nil, err
 	}
 	return true, &KVPair{Key: key, Value: value, LastIndex: uint64(meta.Version)}, nil
@@ -229,12 +232,15 @@ func (s *Zookeeper) AtomicPut(key string, value []byte, previous *KVPair, _ *Wri
 // AtomicDelete deletes a value at "key" if the key has not
 // been modified in the meantime, throws an error if this is the case
 func (s *Zookeeper) AtomicDelete(key string, previous *KVPair) (bool, error) {
-	if previous != nil {
+	if previous == nil {
 		return false, ErrPreviousNotSpecified
 	}
 
 	err := s.client.Delete(normalize(key), int32(previous.LastIndex))
 	if err != nil {
+		if err == zk.ErrBadVersion {
+			return false, ErrKeyModified
+		}
 		return false, err
 	}
 	return true, nil