From 1040ec06e4f83ef6a3b1b50569d232690e930672 Mon Sep 17 00:00:00 2001
From: Alexandre Beslic <abronan@docker.com>
Date: Tue, 26 May 2015 10:39:26 -0700
Subject: [PATCH] Return last index on CAS operation

Signed-off-by: Alexandre Beslic <abronan@docker.com>
---
 pkg/store/consul.go    | 21 +++++++++++++++++----
 pkg/store/etcd.go      |  8 ++++----
 pkg/store/mock.go      |  4 ++--
 pkg/store/store.go     |  2 +-
 pkg/store/zookeeper.go |  4 ++--
 5 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/pkg/store/consul.go b/pkg/store/consul.go
index d0e8e4f729..1fa204bded 100644
--- a/pkg/store/consul.go
+++ b/pkg/store/consul.go
@@ -332,14 +332,27 @@ 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, error) {
+func (s *Consul) AtomicPut(key string, value []byte, previous *KVPair, options *WriteOptions) (bool, *KVPair, error) {
+	lock, err := s.NewLock(key, nil)
+	if err != nil {
+		return false, nil, err
+	}
+
+	lock.Lock()
+	defer lock.Unlock()
+
 	p := &api.KVPair{Key: s.normalize(key), Value: value, ModifyIndex: previous.LastIndex}
 	if work, _, err := s.client.KV().CAS(p, nil); err != nil {
-		return false, err
+		return false, nil, err
 	} else if !work {
-		return false, ErrKeyModified
+		return false, nil, ErrKeyModified
 	}
-	return true, nil
+
+	pair, err := s.Get(key)
+	if err != nil {
+		return false, nil, err
+	}
+	return true, pair, nil
 }
 
 // AtomicDelete deletes a value at "key" if the key has not
diff --git a/pkg/store/etcd.go b/pkg/store/etcd.go
index fb14805902..5fc3dc50c1 100644
--- a/pkg/store/etcd.go
+++ b/pkg/store/etcd.go
@@ -252,12 +252,12 @@ 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, error) {
-	_, err := s.client.CompareAndSwap(normalize(key), string(value), 0, "", previous.LastIndex)
+func (s *Etcd) AtomicPut(key string, value []byte, previous *KVPair, options *WriteOptions) (bool, *KVPair, error) {
+	meta, err := s.client.CompareAndSwap(normalize(key), string(value), 0, "", previous.LastIndex)
 	if err != nil {
-		return false, err
+		return false, nil, err
 	}
-	return true, nil
+	return true, &KVPair{Key: key, Value: value, LastIndex: meta.Node.ModifiedIndex}, nil
 }
 
 // AtomicDelete deletes a value at "key" if the key has not
diff --git a/pkg/store/mock.go b/pkg/store/mock.go
index c5bbcdc431..f7bc22e857 100644
--- a/pkg/store/mock.go
+++ b/pkg/store/mock.go
@@ -75,9 +75,9 @@ func (s *Mock) DeleteTree(prefix string) error {
 }
 
 // AtomicPut mock
-func (s *Mock) AtomicPut(key string, value []byte, previous *KVPair, opts *WriteOptions) (bool, error) {
+func (s *Mock) AtomicPut(key string, value []byte, previous *KVPair, opts *WriteOptions) (bool, *KVPair, error) {
 	args := s.Mock.Called(key, value, previous, opts)
-	return args.Bool(0), args.Error(1)
+	return args.Bool(0), args.Get(1).(*KVPair), args.Error(2)
 }
 
 // AtomicDelete mock
diff --git a/pkg/store/store.go b/pkg/store/store.go
index 0fc390a232..93960ef3ed 100644
--- a/pkg/store/store.go
+++ b/pkg/store/store.go
@@ -87,7 +87,7 @@ type Store interface {
 	DeleteTree(prefix string) error
 
 	// Atomic operation on a single value
-	AtomicPut(key string, value []byte, previous *KVPair, options *WriteOptions) (bool, error)
+	AtomicPut(key string, value []byte, previous *KVPair, options *WriteOptions) (bool, *KVPair, error)
 
 	// Atomic delete of a single value
 	AtomicDelete(key string, previous *KVPair) (bool, error)
diff --git a/pkg/store/zookeeper.go b/pkg/store/zookeeper.go
index 06eebbe1d5..4e611d4dc8 100644
--- a/pkg/store/zookeeper.go
+++ b/pkg/store/zookeeper.go
@@ -214,9 +214,9 @@ 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, options *WriteOptions) (bool, error) {
+func (s *Zookeeper) AtomicPut(key string, value []byte, previous *KVPair, options *WriteOptions) (bool, *KVPair, error) {
 	// Use index of Set method to implement CAS
-	return false, ErrNotImplemented
+	return false, nil, ErrNotImplemented
 }
 
 // AtomicDelete deletes a value at "key" if the key has not