diff --git a/consul/state_store.go b/consul/state_store.go index 7b0808528c..3f1f9a1494 100644 --- a/consul/state_store.go +++ b/consul/state_store.go @@ -782,8 +782,51 @@ func (s *StateStore) KVSDeleteTree() error { } // KVSCheckAndSet is used to perform an atomic check-and-set -func (s *StateStore) KVSCheckAndSet() error { - return nil +func (s *StateStore) KVSCheckAndSet(index uint64, d *structs.DirEntry) (bool, error) { + // Start a new txn + tx, err := s.kvsTable.StartTxn(false, nil) + if err != nil { + return false, err + } + defer tx.Abort() + + // Get the existing node + res, err := s.kvsTable.GetTxn(tx, "id", d.Key) + if err != nil { + return false, err + } + + // Get the existing node if any + var exist *structs.DirEntry + if len(res) > 0 { + exist = res[0].(*structs.DirEntry) + } + + // Use the ModifyIndex as the constraint. A modify of time of 0 + // means we are doing a set-if-not-exists, while any other value + // means we expect that modify time. + if d.ModifyIndex == 0 && exist != nil { + return false, nil + } else if d.ModifyIndex > 0 && (exist == nil || exist.ModifyIndex != d.ModifyIndex) { + return false, nil + } + + // Set the create and modify times + if exist == nil { + d.CreateIndex = index + } else { + d.CreateIndex = exist.CreateIndex + } + d.ModifyIndex = index + + if err := s.kvsTable.InsertTxn(tx, d); err != nil { + return false, err + } + if err := s.kvsTable.SetLastIndexTxn(tx, index); err != nil { + return false, err + } + defer s.watch[s.kvsTable].Notify() + return true, tx.Commit() } // Snapshot is used to create a point in time snapshot diff --git a/consul/state_store_test.go b/consul/state_store_test.go index c2f64dbc7b..c9a26e6763 100644 --- a/consul/state_store_test.go +++ b/consul/state_store_test.go @@ -1044,3 +1044,66 @@ func TestKVSDelete(t *testing.T) { t.Fatalf("bad: %v", d) } } + +func TestKVSCheckAndSet(t *testing.T) { + store, err := testStateStore() + if err != nil { + t.Fatalf("err: %v", err) + } + defer store.Close() + + // CAS should fail, no entry + d := &structs.DirEntry{ + ModifyIndex: 100, + Key: "/foo", + Flags: 42, + Value: []byte("test"), + } + ok, err := store.KVSCheckAndSet(1000, d) + if err != nil { + t.Fatalf("err: %v", err) + } + if ok { + t.Fatalf("unexpected commit") + } + + // Constrain on not-exist, should work + d.ModifyIndex = 0 + ok, err = store.KVSCheckAndSet(1001, d) + if err != nil { + t.Fatalf("err: %v", err) + } + if !ok { + t.Fatalf("expected commit") + } + + // Constrain on not-exist, should fail + d.ModifyIndex = 0 + ok, err = store.KVSCheckAndSet(1002, d) + if err != nil { + t.Fatalf("err: %v", err) + } + if ok { + t.Fatalf("unexpected commit") + } + + // Constrain on a wrong modify time + d.ModifyIndex = 1000 + ok, err = store.KVSCheckAndSet(1003, d) + if err != nil { + t.Fatalf("err: %v", err) + } + if ok { + t.Fatalf("unexpected commit") + } + + // Constrain on a correct modify time + d.ModifyIndex = 1001 + ok, err = store.KVSCheckAndSet(1004, d) + if err != nil { + t.Fatalf("err: %v", err) + } + if !ok { + t.Fatalf("expected commit") + } +}