fix(serving): scan all identifier matches for reply (#5735)

kythe/go/serving/explore/explore_test.go

@@ -453,6 +453,19 @@ func (t protoTable) Lookup(_ context.Context, key []byte, msg proto.Message) err
 	return nil
 }
 
+func (t protoTable) LookupValues(_ context.Context, key []byte, m proto.Message, f func(proto.Message) error) error {
+	val, ok := t[string(key)]
+	if !ok {
+		return nil
+	}
+	msg := m.ProtoReflect().New().Interface()
+	proto.Merge(msg, val)
+	if err := f(msg); err != nil && err != table.ErrStopLookup {
+		return err
+	}
+	return nil
+}
+
 func construct(t *testing.T) *Tables {
 	return &Tables{
 		ParentToChildren:  parentToChildren,

kythe/go/serving/graph/graph_test.go

@@ -546,6 +546,19 @@ func (t testProtoTable) Lookup(_ context.Context, key []byte, msg proto.Message)
 	return nil
 }
 
+func (t testProtoTable) LookupValues(_ context.Context, key []byte, m proto.Message, f func(proto.Message) error) error {
+	val, ok := t[string(key)]
+	if !ok {
+		return nil
+	}
+	msg := m.ProtoReflect().New().Interface()
+	proto.Merge(msg, val)
+	if err := f(msg); err != nil && err != table.ErrStopLookup {
+		return err
+	}
+	return nil
+}
+
 func (t testProtoTable) Buffered() table.BufferedProto { panic("UNIMPLEMENTED") }
 
 func (t testProtoTable) Close(_ context.Context) error { return nil }

kythe/go/serving/identifiers/identifiers.go

@@ -31,6 +31,8 @@ import (
 	"kythe.io/kythe/go/util/kytheuri"
 	"kythe.io/kythe/go/util/log"
 
+	"google.golang.org/protobuf/proto"
+
 	ipb "kythe.io/kythe/proto/identifier_go_proto"
 	srvpb "kythe.io/kythe/proto/serving_go_proto"
 )
@@ -53,31 +55,28 @@ func (it *Table) Find(ctx context.Context, req *ipb.FindRequest) (*ipb.FindReply
 		qname     = req.GetIdentifier()
 		corpora   = req.GetCorpus()
 		languages = req.GetLanguages()
-		match     srvpb.IdentifierMatch
 		reply     ipb.FindReply
 	)
 
-	if err := it.Lookup(ctx, []byte(qname), &match); err != nil {
-		return &reply, nil
-	}
-
-	for _, node := range match.GetNode() {
-		if !validCorpusAndLang(corpora, languages, node) {
-			continue
-		}
-
-		matchNode := ipb.FindReply_Match{
-			Ticket:        node.GetTicket(),
-			NodeKind:      node.GetNodeKind(),
-			NodeSubkind:   node.GetNodeSubkind(),
-			BaseName:      match.GetBaseName(),
-			QualifiedName: match.GetQualifiedName(),
+	return &reply, it.LookupValues(ctx, []byte(qname), (*srvpb.IdentifierMatch)(nil), func(msg proto.Message) error {
+		match := msg.(*srvpb.IdentifierMatch)
+		for _, node := range match.GetNode() {
+			if !validCorpusAndLang(corpora, languages, node) {
+				continue
+			}
+
+			matchNode := &ipb.FindReply_Match{
+				Ticket:        node.GetTicket(),
+				NodeKind:      node.GetNodeKind(),
+				NodeSubkind:   node.GetNodeSubkind(),
+				BaseName:      match.GetBaseName(),
+				QualifiedName: match.GetQualifiedName(),
+			}
+
+			reply.Matches = append(reply.GetMatches(), matchNode)
 		}
-
-		reply.Matches = append(reply.GetMatches(), &matchNode)
-	}
-
-	return &reply, nil
+		return nil
+	})
 }
 
 func validCorpusAndLang(corpora, langs []string, node *srvpb.IdentifierMatch_Node) bool {

kythe/go/serving/identifiers/identifiers_test.go

@@ -20,32 +20,37 @@ import (
 	"context"
 	"testing"
 
-	"google.golang.org/protobuf/proto"
-
 	"kythe.io/kythe/go/storage/table"
 	"kythe.io/kythe/go/test/testutil"
 
+	"google.golang.org/protobuf/proto"
+
 	ipb "kythe.io/kythe/proto/identifier_go_proto"
 	srvpb "kythe.io/kythe/proto/serving_go_proto"
 )
 
 var matchTable = Table{testProtoTable{
-	"foo::bar": &srvpb.IdentifierMatch{
+	"foo::bar": []proto.Message{&srvpb.IdentifierMatch{
 		Node: []*srvpb.IdentifierMatch_Node{
 			node("kythe://corpus?lang=c++", "record", "class"),
+		},
+		BaseName:      "bar",
+		QualifiedName: "foo::bar",
+	}, &srvpb.IdentifierMatch{
+		Node: []*srvpb.IdentifierMatch_Node{
 			node("kythe://corpus?lang=rust", "record", "struct"),
 		},
 		BaseName:      "bar",
 		QualifiedName: "foo::bar",
-	},
+	}},
 
-	"com.java.package.Interface": &srvpb.IdentifierMatch{
+	"com.java.package.Interface": []proto.Message{&srvpb.IdentifierMatch{
 		Node: []*srvpb.IdentifierMatch_Node{
 			node("kythe://habeas?lang=java", "record", "interface"),
 		},
 		BaseName:      "Interface",
 		QualifiedName: "com.java.package.Interface",
-	},
+	}},
 }}
 
 var tests = []testCase{
@@ -118,19 +123,32 @@ type testCase struct {
 	Matches []*ipb.FindReply_Match
 }
 
-type testProtoTable map[string]proto.Message
+type testProtoTable map[string][]proto.Message
 
 func (t testProtoTable) Put(_ context.Context, key []byte, val proto.Message) error {
-	t[string(key)] = val
+	t[string(key)] = []proto.Message{val}
 	return nil
 }
 
 func (t testProtoTable) Lookup(_ context.Context, key []byte, msg proto.Message) error {
 	m, ok := t[string(key)]
-	if !ok {
+	if !ok || len(m) == 0 {
 		return table.ErrNoSuchKey
 	}
-	proto.Merge(msg, m)
+	proto.Merge(msg, m[0])
+	return nil
+}
+
+func (t testProtoTable) LookupValues(_ context.Context, key []byte, m proto.Message, f func(proto.Message) error) error {
+	for _, val := range t[string(key)] {
+		msg := m.ProtoReflect().New().Interface()
+		proto.Merge(msg, val)
+		if err := f(msg); err == table.ErrStopLookup {
+			return nil
+		} else if err != nil {
+			return err
+		}
+	}
 	return nil
 }
 

kythe/go/serving/xrefs/xrefs_test.go

@@ -3550,6 +3550,19 @@ func (t testProtoTable) Lookup(_ context.Context, key []byte, msg proto.Message)
 	return nil
 }
 
+func (t testProtoTable) LookupValues(_ context.Context, key []byte, m proto.Message, f func(proto.Message) error) error {
+	val, ok := t[string(key)]
+	if !ok {
+		return nil
+	}
+	msg := m.ProtoReflect().New().Interface()
+	proto.Merge(msg, val)
+	if err := f(msg); err != nil && err != table.ErrStopLookup {
+		return err
+	}
+	return nil
+}
+
 func (t testProtoTable) Buffered() table.BufferedProto { panic("UNIMPLEMENTED") }
 
 func (t testProtoTable) Close(_ context.Context) error { return nil }

kythe/go/storage/keyvalue/keyvalue.go

@@ -51,6 +51,14 @@ type Range struct {
 	Start, End []byte
 }
 
+// KeyRange returns a Range that contains only the given key.
+func KeyRange(k []byte) *Range {
+	return &Range{
+		Start: k,
+		End:   append(k[0:len(k):len(k)], 0),
+	}
+}
+
 type shard struct {
 	Range
 	count int64

kythe/go/storage/keyvalue/keyvalue_test.go

@@ -26,6 +26,13 @@ import (
 	spb "kythe.io/kythe/proto/storage_go_proto"
 )
 
+func TestKeyRange(t *testing.T) {
+	r := KeyRange([]byte("key"))
+	if c := bytes.Compare(r.Start, r.End); c != -1 {
+		t.Errorf("%s >= %s", r.Start, r.End)
+	}
+}
+
 func TestVNameEncoding(t *testing.T) {
 	tests := []*spb.VName{
 		nil,

kythe/go/storage/table/table.go

@@ -47,6 +47,10 @@ type ProtoLookup interface {
 	// Lookup unmarshals the value for the given key into msg, returning any
 	// error.  If the key was not found, ErrNoSuchKey is returned.
 	Lookup(ctx context.Context, key []byte, msg proto.Message) error
+
+	// Lookup unmarshals the values for the given key into new proto.Message using
+	// m, returning any error.  If the key was not found, f is never called.
+	LookupValues(ctx context.Context, key []byte, m proto.Message, f func(msg proto.Message) error) error
 }
 
 // BufferedProto buffers calls to Put to provide a high throughput write
@@ -65,6 +69,10 @@ type KVProto struct{ keyvalue.DB }
 // ErrNoSuchKey is returned when a value was not found for a particular key.
 var ErrNoSuchKey = errors.New("no such key")
 
+// ErrStopLookup should be returned from the function passed to LookupValues
+// when the client wants no further values.
+var ErrStopLookup = errors.New("stop lookup")
+
 // Lookup implements part of the Proto interface.
 func (t *KVProto) Lookup(ctx context.Context, key []byte, msg proto.Message) error {
 	v, err := t.Get(ctx, key, nil)
@@ -78,6 +86,34 @@ func (t *KVProto) Lookup(ctx context.Context, key []byte, msg proto.Message) err
 	return nil
 }
 
+// LookupValues implements part of the ProtoLookup interface.
+func (t *KVProto) LookupValues(ctx context.Context, key []byte, m proto.Message, f func(proto.Message) error) error {
+	it, err := t.ScanRange(ctx, keyvalue.KeyRange(key), nil)
+	if err != nil {
+		return err
+	}
+	for {
+		_, val, err := it.Next()
+		if err == io.EOF {
+			return nil
+		} else if err != nil {
+			return err
+		}
+
+		msg := m.ProtoReflect().New().Interface()
+		if err := proto.Unmarshal(val, msg); err != nil {
+			return err
+		}
+
+		if err := f(msg); err == ErrStopLookup {
+			return nil
+		} else if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // Put implements part of the Proto interface.
 func (t *KVProto) Put(ctx context.Context, key []byte, msg proto.Message) error {
 	b := t.Buffered()