agent: Adding locking support to KV store

command/agent/kvs_endpoint.go

@@ -156,6 +156,18 @@ func (s *HTTPServer) KVSPut(resp http.ResponseWriter, req *http.Request, args *s
 		applyReq.Op = structs.KVSCAS
 	}
 
+	// Check for lock acquisition
+	if _, ok := params["acquire"]; ok {
+		applyReq.DirEnt.Session = params.Get("acquire")
+		applyReq.Op = structs.KVSLock
+	}
+
+	// Check for lock release
+	if _, ok := params["release"]; ok {
+		applyReq.DirEnt.Session = params.Get("release")
+		applyReq.Op = structs.KVSUnlock
+	}
+
 	// Check the content-length
 	if req.ContentLength > maxKVSize {
 		resp.WriteHeader(413)

command/agent/kvs_endpoint_test.go

@@ -339,3 +339,73 @@ func TestKVSEndpoint_ListKeys(t *testing.T) {
 		}
 	}
 }
+
+func TestKVSEndpoint_AcquireRelease(t *testing.T) {
+	httpTest(t, func(srv *HTTPServer) {
+		// Acquire the lock
+		id := makeTestSession(t, srv)
+		req, err := http.NewRequest("PUT",
+			"/v1/kv/test?acquire="+id, bytes.NewReader(nil))
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+
+		resp := httptest.NewRecorder()
+		obj, err := srv.KVSEndpoint(resp, req)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if res := obj.(bool); !res {
+			t.Fatalf("should work")
+		}
+
+		// Verify we have the lock
+		req, err = http.NewRequest("GET", "/v1/kv/test", nil)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		resp = httptest.NewRecorder()
+		obj, err = srv.KVSEndpoint(resp, req)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		d := obj.(structs.DirEntries)[0]
+
+		// Check the flags
+		if d.Session != id {
+			t.Fatalf("bad: %v", d)
+		}
+
+		// Release the lock
+		req, err = http.NewRequest("PUT",
+			"/v1/kv/test?release="+id, bytes.NewReader(nil))
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		resp = httptest.NewRecorder()
+		obj, err = srv.KVSEndpoint(resp, req)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if res := obj.(bool); !res {
+			t.Fatalf("should work")
+		}
+
+		// Verify we do not have the lock
+		req, err = http.NewRequest("GET", "/v1/kv/test", nil)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		resp = httptest.NewRecorder()
+		obj, err = srv.KVSEndpoint(resp, req)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		d = obj.(structs.DirEntries)[0]
+
+		// Check the flags
+		if d.Session != "" {
+			t.Fatalf("bad: %v", d)
+		}
+	})
+}