Adds in basic query template lookups and vendors newly-updated memdb …

…as well as improved iradix tree.

consul/state/prepared_query.go

@@ -33,18 +33,6 @@ func toPreparedQuery(wrapped interface{}) *structs.PreparedQuery {
 	return wrapped.(*queryWrapper).PreparedQuery
 }
 
-// isQueryWild returns the wild-ness of a query. See isWild for details.
-func isQueryWild(query *structs.PreparedQuery) bool {
-	return query != nil && prepared_query.IsTemplate(query) && query.Name == ""
-}
-
-// isWrappedWild is used to determine if the given wrapped query is a wild one,
-// which means it has an empty Name and it's a template. See the comments for
-// "wild" in schema.go for more details and to see where this is used.
-func isWrappedWild(obj interface{}) (bool, error) {
-	return isQueryWild(toPreparedQuery(obj)), nil
-}
-
 // PreparedQueries is used to pull all the prepared queries from the snapshot.
 func (s *StateSnapshot) PreparedQueries() (structs.PreparedQueries, error) {
 	queries, err := s.tx.Get("prepared-queries", "id")
@@ -123,7 +111,9 @@ func (s *StateStore) preparedQuerySetTxn(tx *memdb.Txn, idx uint64, query *struc
 	}
 
 	// Verify that the query name doesn't already exist, or that we are
-	// updating the same instance that has this name.
+	// updating the same instance that has this name. If this is a template
+	// and the name is empty then we make sure there's not an empty template
+	// already registered.
 	if query.Name != "" {
 		wrapped, err := tx.First("prepared-queries", "name", query.Name)
 		if err != nil {
@@ -133,18 +123,14 @@ func (s *StateStore) preparedQuerySetTxn(tx *memdb.Txn, idx uint64, query *struc
 		if other != nil && (existing == nil || existing.ID != other.ID) {
 			return fmt.Errorf("name '%s' aliases an existing query name", query.Name)
 		}
-	}
-
-	// Similarly, if this is the wild query make sure there isn't another
-	// one, or that we are updating the same one.
-	if isQueryWild(query) {
-		wrapped, err := tx.First("prepared-queries", "wild", true)
+	} else if prepared_query.IsTemplate(query) {
+		wrapped, err := tx.First("prepared-queries", "template", query.Name)
 		if err != nil {
 			return fmt.Errorf("failed prepared query lookup: %s", err)
 		}
 		other := toPreparedQuery(wrapped)
 		if other != nil && (existing == nil || existing.ID != other.ID) {
-			return fmt.Errorf("a prepared query template already exists with an empty name")
+			return fmt.Errorf("name '%s' aliases an existing query template name", query.Name)
 		}
 	}
 
@@ -311,27 +297,22 @@ func (s *StateStore) PreparedQueryResolve(queryIDOrName string) (uint64, *struct
 		}
 	}
 
-	// Then try by name. We use a prefix match but check to make sure that
-	// the query's name matches the whole prefix for a non-template query.
-	// Templates are allowed to use the partial match. It's more efficient
-	// to combine the two lookups here, even though the logic is a little
-	// less clear.
+	// Next, look for an exact name match. This is the common case for static
+	// prepared queries, and could also apply to templates.
 	{
-		wrapped, err := tx.First("prepared-queries", "name_prefix", queryIDOrName)
+		wrapped, err := tx.First("prepared-queries", "name", queryIDOrName)
 		if err != nil {
 			return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
 		}
 		if wrapped != nil {
-			query := toPreparedQuery(wrapped)
-			if query.Name == queryIDOrName || prepared_query.IsTemplate(query) {
-				return prep(wrapped)
-			}
+			return prep(wrapped)
 		}
 	}
 
-	// Finally, see if there's a wild template we can use.
+	// Next, look for the longest prefix match among the prepared query
+	// templates.
 	{
-		wrapped, err := tx.First("prepared-queries", "wild", true)
+		wrapped, err := tx.LongestPrefix("prepared-queries", "template_prefix", queryIDOrName)
 		if err != nil {
 			return 0, nil, fmt.Errorf("failed prepared query lookup: %s", err)
 		}

consul/state/prepared_query_index.go

@@ -0,0 +1,51 @@
+package state
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/hashicorp/consul/consul/prepared_query"
+)
+
+// PreparedQueryIndex is a custom memdb indexer used to manage index prepared
+// query templates. None of the built-in indexers do what we need, and our
+// use case is pretty specific so it's better to put the logic here.
+type PreparedQueryIndex struct {
+}
+
+// FromObject is used to compute the index key when inserting or updating an
+// object.
+func (*PreparedQueryIndex) FromObject(obj interface{}) (bool, []byte, error) {
+	wrapped, ok := obj.(*queryWrapper)
+	if !ok {
+		return false, nil, fmt.Errorf("invalid object given to index as prepared query")
+	}
+
+	query := toPreparedQuery(wrapped)
+	if !prepared_query.IsTemplate(query) {
+		return false, nil, nil
+	}
+
+	// Always prepend a null so that we can represent even an empty name.
+	out := "\x00" + strings.ToLower(query.Name)
+	return true, []byte(out), nil
+}
+
+// FromArgs is used when querying for an exact match. Since we don't add any
+// suffix we can just call the prefix version.
+func (p *PreparedQueryIndex) FromArgs(args ...interface{}) ([]byte, error) {
+	return p.PrefixFromArgs(args...)
+}
+
+// PrefixFromArgs is used when doing a prefix scan for an object.
+func (*PreparedQueryIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
+	if len(args) != 1 {
+		return nil, fmt.Errorf("must provide only a single argument")
+	}
+	arg, ok := args[0].(string)
+	if !ok {
+		return nil, fmt.Errorf("argument must be a string: %#v", args[0])
+	}
+	arg = "\x00" + strings.ToLower(arg)
+	return []byte(arg), nil
+}

consul/state/prepared_query_test.go

@@ -525,9 +525,12 @@ func TestStateStore_PreparedQuery_Snapshot_Restore(t *testing.T) {
 		},
 		&structs.PreparedQuery{
 			ID:   testUUID(),
-			Name: "bob",
+			Name: "bob-",
+			Template: structs.QueryTemplateOptions{
+				Type: structs.QueryTemplateTypeNamePrefixMatch,
+			},
 			Service: structs.ServiceQuery{
-				Service: "mongodb",
+				Service: "${name.suffix}",
 			},
 		},
 	}
@@ -571,9 +574,12 @@ func TestStateStore_PreparedQuery_Snapshot_Restore(t *testing.T) {
 		},
 		&structs.PreparedQuery{
 			ID:   queries[1].ID,
-			Name: "bob",
+			Name: "bob-",
+			Template: structs.QueryTemplateOptions{
+				Type: structs.QueryTemplateTypeNamePrefixMatch,
+			},
 			Service: structs.ServiceQuery{
-				Service: "mongodb",
+				Service: "${name.suffix}",
 			},
 			RaftIndex: structs.RaftIndex{
 				CreateIndex: 5,
@@ -612,6 +618,19 @@ func TestStateStore_PreparedQuery_Snapshot_Restore(t *testing.T) {
 		if !reflect.DeepEqual(actual, expected) {
 			t.Fatalf("bad: %v", actual)
 		}
+
+		// Make sure the second query, which is a template, was compiled
+		// and can be resolved.
+		_, query, err := s.PreparedQueryResolve("bob-backwards-is-bob")
+		if err != nil {
+			t.Fatalf("err: %s", err)
+		}
+		if query == nil {
+			t.Fatalf("should have resolved the query")
+		}
+		if query.Service.Service != "backwards-is-bob" {
+			t.Fatalf("bad: %s", query.Service.Service)
+		}
 	}()
 }
 

consul/state/schema.go

@@ -390,20 +390,11 @@ func preparedQueriesTableSchema() *memdb.TableSchema {
 					Lowercase: true,
 				},
 			},
-			// This is a bit of an oddball. It's an important feature
-			// of prepared query templates to be able to define a
-			// single template that matches any query. Unfortunately,
-			// we can't index an empty Name field. This index lets us
-			// keep track of whether there is any wild template in
-			// existence, so there will be one "true" in here if that
-			// exists, and everything else will be "false".
-			"wild": &memdb.IndexSchema{
-				Name:         "wild",
-				AllowMissing: false,
-				Unique:       false,
-				Indexer: &memdb.ConditionalIndex{
-					Conditional: isWrappedWild,
-				},
+			"template": &memdb.IndexSchema{
+				Name:         "template",
+				AllowMissing: true,
+				Unique:       true,
+				Indexer:      &PreparedQueryIndex{},
 			},
 			"session": &memdb.IndexSchema{
 				Name:         "session",