Fixing bug in state store

consul/catalog_endpoint_test.go

@@ -337,3 +337,49 @@ func TestCatalogNodeServices(t *testing.T) {
 		t.Fatalf("bad: %v", out)
 	}
 }
+
+// Used to check for a regression against a known bug
+func TestCatalogRegister_FailedCase1(t *testing.T) {
+	dir1, s1 := testServer(t)
+	defer os.RemoveAll(dir1)
+	defer s1.Shutdown()
+	client := rpcClient(t, s1)
+	defer client.Close()
+
+	arg := structs.RegisterRequest{
+		Datacenter:  "dc1",
+		Node:        "bar",
+		Address:     "127.0.0.2",
+		ServiceName: "web",
+		ServiceTag:  "",
+		ServicePort: 8000,
+	}
+	var out struct{}
+
+	err := client.Call("Catalog.Register", &arg, &out)
+	if err == nil || err.Error() != "No cluster leader" {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Wait for leader
+	time.Sleep(100 * time.Millisecond)
+
+	if err := client.Call("Catalog.Register", &arg, &out); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Check we can get this back
+	query := &structs.ServiceNodesRequest{
+		Datacenter:  "dc1",
+		ServiceName: "web",
+	}
+	var nodes structs.ServiceNodes
+	if err := client.Call("Catalog.ServiceNodes", query, &nodes); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Check the output
+	if len(nodes) != 1 {
+		t.Fatalf("Bad: %v", nodes)
+	}
+}

consul/state_store.go

@@ -463,8 +463,16 @@ func parseServiceNodes(tx *mdb.Txn, index mdb.DBI, prefix []byte) structs.Servic
 
 	var nodes structs.ServiceNodes
 	var node structs.ServiceNode
+	var key, val []byte
+	first := true
+
 	for {
-		key, val, err := cursor.Get(nil, mdb.NEXT)
+		if first {
+			first = false
+			key, val, err = cursor.Get(prefix, mdb.SET_RANGE)
+		} else {
+			key, val, err = cursor.Get(nil, mdb.NEXT)
+		}
 		if err == mdb.NotFound {
 			break
 		} else if err != nil {

consul/state_store_test.go

@@ -216,6 +216,14 @@ func TestServiceNodes(t *testing.T) {
 		t.Fatalf("err: %v")
 	}
 
+	if err := store.EnsureService("foo", "api", "", 5000); err != nil {
+		t.Fatalf("err: %v")
+	}
+
+	if err := store.EnsureService("bar", "api", "", 5000); err != nil {
+		t.Fatalf("err: %v")
+	}
+
 	if err := store.EnsureService("foo", "db", "master", 8000); err != nil {
 		t.Fatalf("err: %v")
 	}