From 473c0796970232134b713a01971ff8064b90728d Mon Sep 17 00:00:00 2001 From: Nia Weiss Date: Mon, 12 Oct 2020 10:59:40 -0400 Subject: [PATCH] apply "maybe" key translation WIP --- api.go | 32 ++++ boltdb/translate.go | 69 +++++++++ boltdb/translate_test.go | 118 +++++++++++++++ client.go | 39 +++-- cluster.go | 316 +++++++++++++++++++++++++++++++++++++++ http/client.go | 202 +++++++++++++++++++++++++ http/handler.go | 185 +++++++++++++++++++++++ mock/translator.go | 12 +- translate.go | 62 +++++++- 9 files changed, 1019 insertions(+), 16 deletions(-) diff --git a/api.go b/api.go index 7c11a4e3e..88a1e6be0 100644 --- a/api.go +++ b/api.go @@ -1661,6 +1661,38 @@ func (api *API) TranslateIDs(ctx context.Context, r io.Reader) (_ []byte, err er return buf, nil } +// FindIndexKeys looks up column keys in the index, mapping them to IDs. +// If a key does not exist, it will be absent from the resulting map. +func (api *API) FindIndexKeys(ctx context.Context, index string, keys ...string) (map[string]uint64, error) { + return api.cluster.findIndexKeys(ctx, index, keys...) +} + +// FindFieldKeys looks up keys in a field, mapping them to IDs. +// If a key does not exist, it will be absent from the resulting map. +func (api *API) FindFieldKeys(ctx context.Context, index, field string, keys ...string) (map[string]uint64, error) { + f := api.holder.Field(index, field) + if f == nil { + return nil, newNotFoundError(ErrFieldNotFound, field) + } + return api.cluster.findFieldKeys(ctx, f, keys...) +} + +// CreateIndexKeys looks up column keys in the index, mapping them to IDs. +// If a key does not exist, it will be created. +func (api *API) CreateIndexKeys(ctx context.Context, index string, keys ...string) (map[string]uint64, error) { + return api.cluster.createIndexKeys(ctx, index, keys...) +} + +// CreateFieldKeys looks up keys in a field, mapping them to IDs. +// If a key does not exist, it will be created. +func (api *API) CreateFieldKeys(ctx context.Context, index, field string, keys ...string) (map[string]uint64, error) { + f := api.holder.Field(index, field) + if f == nil { + return nil, newNotFoundError(ErrFieldNotFound, field) + } + return api.cluster.createFieldKeys(ctx, f, keys...) +} + // PrimaryReplicaNodeURL returns the URL of the cluster's primary replica. func (api *API) PrimaryReplicaNodeURL() url.URL { node := api.cluster.PrimaryReplicaNode() diff --git a/boltdb/translate.go b/boltdb/translate.go index d37dae0d9..46bb8f8c9 100644 --- a/boltdb/translate.go +++ b/boltdb/translate.go @@ -190,6 +190,75 @@ func (s *TranslateStore) TranslateKeys(keys []string, writable bool) ([]uint64, return s.translateKeys(keys, writable) } +// FindKeys looks up the ID for each key. +// Keys are not created if they do not exist. +// Missing keys are not considered errors, so the length of the result may be less than that of the input. +func (s *TranslateStore) FindKeys(keys ...string) (map[string]uint64, error) { + result := make(map[string]uint64, len(keys)) + err := s.db.View(func(tx *bolt.Tx) error { + bkt := tx.Bucket(bucketKeys) + if bkt == nil { + return errors.Errorf(errFmtTranslateBucketNotFound, bucketKeys) + } + for _, key := range keys { + id, _ := findIDByKey(bkt, key) + if id == 0 { + // The key does not exist. + continue + } + + result[key] = id + } + return nil + }) + if err != nil { + return nil, err + } + + return result, nil +} + +// CreateKeys maps all keys to IDs, creating the IDs if they do not exist. +// If the translator is read-only, this will return an error. +func (s *TranslateStore) CreateKeys(keys ...string) (map[string]uint64, error) { + if s.ReadOnly() { + return nil, pilosa.ErrTranslateStoreReadOnly + } + + written := false + result := make(map[string]uint64, len(keys)) + err := s.db.Update(func(tx *bolt.Tx) error { + bkt := tx.Bucket(bucketKeys) + if bkt == nil { + return errors.Errorf(errFmtTranslateBucketNotFound, bucketKeys) + } + for _, key := range keys { + id, boltKey := findIDByKey(bkt, key) + if id == 0 { + // The key does not exist. + id = pilosa.GenerateNextPartitionedID(s.index, maxID(tx), s.partitionID, s.partitionN) + if err := bkt.Put(boltKey, u64tob(id)); err != nil { + return err + } else if err := tx.Bucket(bucketIDs).Put(u64tob(id), boltKey); err != nil { + return err + } + written = true + } + + result[key] = id + } + return nil + }) + if err != nil { + return nil, err + } + if written { + s.notifyWrite() + } + + return result, nil +} + func (s *TranslateStore) translateKeys(keys []string, writable bool) ([]uint64, error) { ids := make([]uint64, 0, len(keys)) diff --git a/boltdb/translate_test.go b/boltdb/translate_test.go index 73b4e5345..271325e94 100644 --- a/boltdb/translate_test.go +++ b/boltdb/translate_test.go @@ -92,6 +92,46 @@ func TestTranslateStore_TranslateKeys(t *testing.T) { } } +func TestTranslateStore_CreateKeys(t *testing.T) { + s := MustOpenNewTranslateStore() + defer MustCloseTranslateStore(s) + + ids, err := s.CreateKeys("abc", "abc") + if err != nil { + t.Fatal(err) + } else if _, ok := ids["abc"]; !ok { + t.Fatalf(`missing "abc"; got %v`, ids) + } else if len(ids) > 1 { + t.Fatalf("expected one key, got %d in %v", len(ids), ids) + } + + // Ensure different keys translate to different IDs. + ids1, err := s.CreateKeys("foo", "bar") + if err != nil { + t.Fatal(err) + } else if foo, bar := ids1["foo"], ids1["bar"]; foo == bar { + t.Fatalf(`"foo" and "bar" map back to the same ID %d`, foo) + } + + // Ensure retranslation returns original IDs. + if ids, err := s.CreateKeys("bar", "foo"); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(ids, ids1) { + t.Fatalf("retranslation produced result %v which is different from original translation %v", ids, ids1) + } + + // Ensure retranslating with existing and non-existing keys returns correctly. + if ids, err := s.CreateKeys("foo", "baz", "bar"); err != nil { + t.Fatal(err) + } else if got, want := ids["foo"], ids1["foo"]; got != want { + t.Fatalf(`mismatched ID %d for "foo" (previously %d)`, got, want) + } else if _, ok := ids["baz"]; !ok { + t.Fatalf(`missing translation for "baz"; got %v`, ids) + } else if got, want := ids["bar"], ids1["bar"]; got != want { + t.Fatalf(`mismatched ID %d for "bar" (previously %d)`, got, want) + } +} + func TestTranslateStore_ReadKey(t *testing.T) { s := MustOpenNewTranslateStore() defer MustCloseTranslateStore(s) @@ -216,6 +256,84 @@ func TestTranslateStore_TranslateIDs(t *testing.T) { } } +func TestTranslateStore_FindKeys(t *testing.T) { + cases := []struct { + name string + data []string + lookup []string + }{ + { + name: "All", + data: []string{"plugh", "xyzzy", "h"}, + lookup: []string{"plugh", "xyzzy", "h"}, + }, + { + name: "Extra", + data: []string{"plugh", "xyzzy", "h"}, + lookup: []string{"plugh", "xyzzy", "h", "65"}, + }, + { + name: "None", + data: []string{"a", "b", "c"}, + lookup: []string{"d", "e"}, + }, + { + name: "Empty", + lookup: []string{"h"}, + }, + { + name: "LookupNothing", + }, + } + + for _, c := range cases { + c := c + t.Run(c.name, func(t *testing.T) { + s := MustOpenNewTranslateStore() + defer MustCloseTranslateStore(s) + + var naieveMap map[string]uint64 + if c.data != nil { + // Load in key data. + keys := c.data + ids, err := s.TranslateKeys(keys, true) + if err != nil { + t.Errorf("failed to import keys: %v", err) + return + } + if len(ids) != len(keys) { + t.Errorf("mapped %d keys to %d ids", len(keys), len(ids)) + return + } + naieveMap = make(map[string]uint64, len(keys)) + for i, key := range keys { + naieveMap[key] = ids[i] + } + } + + // Compute expected lookup result. + result := map[string]uint64{} + for _, key := range c.lookup { + id, ok := naieveMap[key] + if !ok { + // The key is expected to be missing. + continue + } + + result[key] = id + } + + // Find the keys. + found, err := s.FindKeys(c.lookup...) + if err != nil { + t.Errorf("failed to find keys: %v", err) + } else if !reflect.DeepEqual(result, found) { + t.Errorf("expected %v but found %v", result, found) + } + }) + } +} + func TestTranslateStore_EntryReader(t *testing.T) { t.Run("OK", func(t *testing.T) { s := MustOpenNewTranslateStore() diff --git a/client.go b/client.go index 42ec51ba0..28e79679c 100644 --- a/client.go +++ b/client.go @@ -92,11 +92,17 @@ type InternalQueryClient interface { // Trasnlate keys on the particular node. The parameter writable informs TranslateStore if we can generate a new ID if any of keys does not exist. TranslateKeysNode(ctx context.Context, uri *URI, index, field string, keys []string, writable bool) ([]uint64, error) TranslateIDsNode(ctx context.Context, uri *URI, index, field string, id []uint64) ([]string, error) + + FindIndexKeysNode(ctx context.Context, uri *URI, index string, keys ...string) (map[string]uint64, error) + FindFieldKeysNode(ctx context.Context, uri *URI, index string, field string, keys ...string) (map[string]uint64, error) + + CreateIndexKeysNode(ctx context.Context, uri *URI, index string, keys ...string) (map[string]uint64, error) + CreateFieldKeysNode(ctx context.Context, uri *URI, index string, field string, keys ...string) (map[string]uint64, error) } type nopInternalQueryClient struct{} -func (n *nopInternalQueryClient) QueryNode(ctx context.Context, uri *URI, index string, queryRequest *QueryRequest) (*QueryResponse, error) { +func (n nopInternalQueryClient) QueryNode(ctx context.Context, uri *URI, index string, queryRequest *QueryRequest) (*QueryResponse, error) { return nil, nil } @@ -108,15 +114,31 @@ func (n nopInternalQueryClient) TranslateIDsNode(ctx context.Context, uri *URI, return nil, nil } -func newNopInternalQueryClient() *nopInternalQueryClient { - return &nopInternalQueryClient{} +func (n nopInternalQueryClient) FindIndexKeysNode(ctx context.Context, uri *URI, index string, keys ...string) (map[string]uint64, error) { + return nil, nil +} + +func (n nopInternalQueryClient) FindFieldKeysNode(ctx context.Context, uri *URI, index string, field string, keys ...string) (map[string]uint64, error) { + return nil, nil +} + +func (n nopInternalQueryClient) CreateIndexKeysNode(ctx context.Context, uri *URI, index string, keys ...string) (map[string]uint64, error) { + return nil, nil +} + +func (n nopInternalQueryClient) CreateFieldKeysNode(ctx context.Context, uri *URI, index string, field string, keys ...string) (map[string]uint64, error) { + return nil, nil +} + +func newNopInternalQueryClient() nopInternalQueryClient { + return nopInternalQueryClient{} } var _ InternalQueryClient = newNopInternalQueryClient() //=============== -type nopInternalClient struct{} +type nopInternalClient struct{ nopInternalQueryClient } func newNopInternalClient() nopInternalClient { return nopInternalClient{} @@ -144,15 +166,6 @@ func (n nopInternalClient) Nodes(ctx context.Context) ([]*Node, error) { func (n nopInternalClient) Query(ctx context.Context, index string, queryRequest *QueryRequest) (*QueryResponse, error) { return nil, nil } -func (n nopInternalClient) QueryNode(ctx context.Context, uri *URI, index string, queryRequest *QueryRequest) (*QueryResponse, error) { - return nil, nil -} -func (n nopInternalClient) TranslateKeysNode(ctx context.Context, uri *URI, index, field string, keys []string, writable bool) ([]uint64, error) { - return nil, nil -} -func (n nopInternalClient) TranslateIDsNode(ctx context.Context, uri *URI, index, field string, ids []uint64) ([]string, error) { - return nil, nil -} func (n nopInternalClient) Import(ctx context.Context, index, field string, shard uint64, bits []Bit, opts ...ImportOption) error { return nil } diff --git a/cluster.go b/cluster.go index 3905c91d1..0296bd487 100644 --- a/cluster.go +++ b/cluster.go @@ -2359,6 +2359,149 @@ func (c *cluster) translateFieldKeys(ctx context.Context, field *Field, keys []s return ids, nil } +func (c *cluster) findFieldKeys(ctx context.Context, field *Field, keys ...string) (map[string]uint64, error) { + // Attempt to find the keys locally. + localTranslations, err := field.TranslateStore().FindKeys(keys...) + if err != nil { + return nil, errors.Wrapf(err, "translating field(%s/%s) keys(%v) locally", field.Index(), field.Name(), keys) + } + + // Check for missing keys. + var missing []string + if len(keys) > len(localTranslations) { + // There are either duplicate keys or missing keys. + // This should work either way. + missing = make([]string, 0, len(keys)-len(localTranslations)) + for _, k := range keys { + _, found := localTranslations[k] + if !found { + missing = append(missing, k) + } + } + } + if len(missing) == 0 { + // All keys were available locally. + return localTranslations, nil + } + + // It is possible that the missing keys exist, but have not been synced to the local replica. + coordinator := c.coordinatorNode() + if coordinator == nil { + return nil, errors.Errorf("translating field(%s/%s) keys(%v) - cannot find coordinator node", field.Index(), field.Name(), keys) + } + if c.Node.ID == coordinator.ID { + // The local copy is the authoritative copy. + return localTranslations, nil + } + + // Forward the missing keys to the coordinator. + // The coordinator has the authoritative copy. + remoteTranslations, err := c.InternalClient.FindFieldKeysNode(ctx, &coordinator.URI, field.Index(), field.Name(), keys...) + if err != nil { + return nil, errors.Wrapf(err, "translating field(%s/%s) keys(%v) remotely", field.Index(), field.Name(), keys) + } + + // Merge the remote translations into the local translations. + translations := localTranslations + for key, id := range remoteTranslations { + translations[key] = id + } + + return translations, nil +} + +func (c *cluster) createFieldKeys(ctx context.Context, field *Field, keys ...string) (map[string]uint64, error) { + // The coordinator is the only node that can create field keys, since it owns the authoritative copy. + coordinator := c.coordinatorNode() + if coordinator == nil { + return nil, errors.Errorf("translating field(%s/%s) keys(%v) - cannot find coordinator node", field.Index(), field.Name(), keys) + } + if c.Node.ID == coordinator.ID { + // The local copy is the authoritative copy. + return field.TranslateStore().CreateKeys(keys...) + } + + // Attempt to find the keys locally. + // They cannot be created locally, but skipping keys that exist can reduce network usage. + localTranslations, err := field.TranslateStore().FindKeys(keys...) + if err != nil { + return nil, errors.Wrapf(err, "translating field(%s/%s) keys(%v) locally", field.Index(), field.Name(), keys) + } + + // Check for missing keys. + var missing []string + if len(keys) > len(localTranslations) { + // There are either duplicate keys or missing keys. + // This should work either way. + missing = make([]string, 0, len(keys)-len(localTranslations)) + for _, k := range keys { + _, found := localTranslations[k] + if !found { + missing = append(missing, k) + } + } + } + if len(missing) == 0 { + // All keys exist locally. + // There is no need to create anything. + return localTranslations, nil + } + + // Forward the missing keys to the coordinator to be created. + remoteTranslations, err := c.InternalClient.CreateFieldKeysNode(ctx, &coordinator.URI, field.Index(), field.Name(), keys...) + if err != nil { + return nil, errors.Wrapf(err, "translating field(%s/%s) keys(%v) remotely", field.Index(), field.Name(), keys) + } + + // Merge the remote translations into the local translations. + translations := localTranslations + for key, id := range remoteTranslations { + translations[key] = id + } + + return translations, nil +} + +func (c *cluster) translateFieldIDs(field *Field, ids map[uint64]struct{}) (map[uint64]string, error) { + idList := make([]uint64, len(ids)) + { + i := 0 + for id := range ids { + idList[i] = id + i++ + } + } + + keyList, err := c.translateFieldListIDs(field, idList) + if err != nil { + return nil, err + } + + mapped := make(map[uint64]string, len(idList)) + for i, key := range keyList { + mapped[idList[i]] = key + } + return mapped, nil +} + +func (c *cluster) translateFieldListIDs(field *Field, ids []uint64) (keys []string, err error) { + coordinator := c.coordinatorNode() + if coordinator == nil { + return nil, errors.Errorf("translating field(%s/%s) ids(%v) - cannot find coordinator node", field.Index(), field.Name(), ids) + } + + if c.Node.ID == coordinator.ID { + keys, err = field.TranslateStore().TranslateIDs(ids) + } else { + keys, err = c.InternalClient.TranslateIDsNode(context.Background(), &coordinator.URI, field.Index(), field.Name(), ids) + } + if err != nil { + return nil, errors.Wrapf(err, "translating field(%s/%s) ids(%v)", field.Index(), field.Name(), ids) + } + + return keys, err +} + func (c *cluster) translateIndexKey(ctx context.Context, indexName string, key string, writable bool) (uint64, error) { keyMap, err := c.translateIndexKeySet(ctx, indexName, map[string]struct{}{key: struct{}{}}, writable) if err != nil { @@ -2450,6 +2593,179 @@ func (c *cluster) translateIndexKeySet(ctx context.Context, indexName string, ke return keyMap, nil } +func (c *cluster) findIndexKeys(ctx context.Context, indexName string, keys ...string) (map[string]uint64, error) { + idx := c.holder.Index(indexName) + if idx == nil { + return nil, ErrIndexNotFound + } + + // Split keys by partition. + keysByPartition := make(map[int][]string, c.partitionN) + for _, key := range keys { + partitionID := c.Topology.KeyPartition(indexName, key) + keysByPartition[partitionID] = append(keysByPartition[partitionID], key) + } + + // TODO: use local replicas to short-circuit network traffic + + // Group keys by node. + // Delete remote keys from the by-partition map so that it can be used for local translation. + keysByNode := make(map[*Node][]string) + for partitionID, keys := range keysByPartition { + // Find the primary node for this partition. + primary := c.primaryPartitionNode(partitionID) + if primary == nil { + return nil, errors.Errorf("translating index(%s) keys(%v) on partition(%d) - cannot find primary node", indexName, keys, partitionID) + } + + if c.Node.ID == primary.ID { + // The partition is local. + continue + } + + // Group the partition to be processed remotely. + keysByNode[primary] = append(keysByNode[primary], keys...) + delete(keysByPartition, partitionID) + } + + // Start translating keys remotely. + // On child calls, there are no remote results. + remoteResults := make(chan map[string]uint64, len(keysByNode)) + var g errgroup.Group + defer g.Wait() + for node, keys := range keysByNode { + node, keys := node, keys + + g.Go(func() error { + translations, err := c.InternalClient.FindIndexKeysNode(ctx, &node.URI, indexName, keys...) + if err != nil { + return errors.Wrapf(err, "translating index(%s) keys(%v) on node %s", indexName, keys, node.ID) + } + + remoteResults <- translations + return nil + }) + } + + // Translate local keys. + translations := make(map[string]uint64) + for partitionID, keys := range keysByPartition { + // Find the keys within the partition. + t, err := idx.TranslateStore(partitionID).FindKeys(keys...) + if err != nil { + return nil, errors.Wrapf(err, "translating index(%s) keys(%v) on partition(%d)", idx.Name(), keys, partitionID) + } + + // Merge the translations from this partition. + for key, id := range t { + translations[key] = id + } + } + + // Wait for remote key sets. + if err := g.Wait(); err != nil { + return nil, err + } + + // Merge the translations. + close(remoteResults) + for t := range remoteResults { + for key, id := range t { + translations[key] = id + } + } + return translations, nil +} + +func (c *cluster) createIndexKeys(ctx context.Context, indexName string, keys ...string) (map[string]uint64, error) { + idx := c.holder.Index(indexName) + if idx == nil { + return nil, ErrIndexNotFound + } + + // Split keys by partition. + keysByPartition := make(map[int][]string, c.partitionN) + for _, key := range keys { + partitionID := c.Topology.KeyPartition(indexName, key) + keysByPartition[partitionID] = append(keysByPartition[partitionID], key) + } + + // TODO: use local replicas to short-circuit network traffic + + // Group keys by node. + // Delete remote keys from the by-partition map so that it can be used for local translation. + keysByNode := make(map[*Node][]string) + for partitionID, keys := range keysByPartition { + // Find the primary node for this partition. + primary := c.primaryPartitionNode(partitionID) + if primary == nil { + return nil, errors.Errorf("translating index(%s) keys(%v) on partition(%d) - cannot find primary node", indexName, keys, partitionID) + } + + if c.Node.ID == primary.ID { + // The partition is local. + continue + } + + // Group the partition to be processed remotely. + keysByNode[primary] = append(keysByNode[primary], keys...) + delete(keysByPartition, partitionID) + } + + translateResults := make(chan map[string]uint64, len(keysByNode)+len(keysByPartition)) + var g errgroup.Group + defer g.Wait() + + // Start translating keys remotely. + // On child calls, there are no remote results. + for node, keys := range keysByNode { + node, keys := node, keys + + g.Go(func() error { + translations, err := c.InternalClient.CreateIndexKeysNode(ctx, &node.URI, indexName, keys...) + if err != nil { + return errors.Wrapf(err, "translating index(%s) keys(%v) on node %s", indexName, keys, node.ID) + } + + translateResults <- translations + return nil + }) + } + + // Translate local keys. + // TODO: make this less horrible (why fsync why?????) + // This is kinda terrible because each goroutine does an fsync, thus locking up an entire OS thread. + // AHHHHHHHHHHHHHHHHHH + for partitionID, keys := range keysByPartition { + partitionID, keys := partitionID, keys + + g.Go(func() error { + translations, err := idx.TranslateStore(partitionID).CreateKeys(keys...) + if err != nil { + return errors.Wrapf(err, "translating index(%s) keys(%v) on partition(%d)", idx.Name(), keys, partitionID) + } + + translateResults <- translations + return nil + }) + } + + // Wait for remote key sets. + if err := g.Wait(); err != nil { + return nil, err + } + + // Merge the translations. + translations := make(map[string]uint64, len(keys)) + close(translateResults) + for t := range translateResults { + for key, id := range t { + translations[key] = id + } + } + return translations, nil +} + func (c *cluster) translateIndexIDs(ctx context.Context, indexName string, ids []uint64) ([]string, error) { idSet := make(map[uint64]struct{}) for _, id := range ids { diff --git a/http/client.go b/http/client.go index 986dd6afe..dab366a2c 100644 --- a/http/client.go +++ b/http/client.go @@ -1247,6 +1247,208 @@ func (c *InternalClient) TranslateIDsNode(ctx context.Context, uri *pilosa.URI, return tkresp.Keys, nil } +func (c *InternalClient) FindIndexKeysNode(ctx context.Context, uri *pilosa.URI, index string, keys ...string) (transMap map[string]uint64, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.FindKeysNodePartition") + defer span.Finish() + + // Create HTTP request. + u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/index/%s/keys", index)) + reqData, err := json.Marshal(keys) + if err != nil { + return nil, errors.Wrap(err, "marshalling request") + } + req, err := http.NewRequest("POST", u.String(), bytes.NewReader(reqData)) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + // Apply headers. + req.Header.Set("Content-Length", strconv.Itoa(len(reqData))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + + // Send the request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, errors.Wrap(err, "executing request") + } + defer func() { + cerr := resp.Body.Close() + if cerr != nil && err == nil { + err = errors.Wrap(cerr, "closing response body") + } + }() + + // Read the response body. + result, err := ioutil.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading response") + } + + // Decode the translations. + transMap = make(map[string]uint64, len(keys)) + err = json.Unmarshal(result, &transMap) + if err != nil { + return nil, errors.Wrap(err, "json decoding") + } + + return transMap, nil +} + +func (c *InternalClient) FindFieldKeysNode(ctx context.Context, uri *pilosa.URI, index string, field string, keys ...string) (transMap map[string]uint64, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.FindKeysFieldNode") + defer span.Finish() + + // Create HTTP request. + u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/index/%s/field/%s", index, field)) + q := u.Query() + q.Add("remote", "true") + u.RawQuery = q.Encode() + reqData, err := json.Marshal(keys) + if err != nil { + return nil, errors.Wrap(err, "marshalling request") + } + req, err := http.NewRequest("POST", u.String(), bytes.NewReader(reqData)) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + // Apply headers. + req.Header.Set("Content-Length", strconv.Itoa(len(reqData))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + + // Send the request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, errors.Wrap(err, "executing request") + } + defer func() { + cerr := resp.Body.Close() + if cerr != nil && err == nil { + err = errors.Wrap(cerr, "closing response body") + } + }() + + // Read the response body. + result, err := ioutil.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading response") + } + + // Decode the translations. + transMap = make(map[string]uint64, len(keys)) + err = json.Unmarshal(result, &transMap) + if err != nil { + return nil, errors.Wrap(err, "json decoding") + } + + return transMap, nil +} + +func (c *InternalClient) CreateIndexKeysNode(ctx context.Context, uri *pilosa.URI, index string, keys ...string) (transMap map[string]uint64, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.FindKeysNodePartition") + defer span.Finish() + + // Create HTTP request. + u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/index/%s/keys", index)) + reqData, err := json.Marshal(keys) + if err != nil { + return nil, errors.Wrap(err, "marshalling request") + } + req, err := http.NewRequest("PUT", u.String(), bytes.NewReader(reqData)) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + // Apply headers. + req.Header.Set("Content-Length", strconv.Itoa(len(reqData))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + + // Send the request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, errors.Wrap(err, "executing request") + } + defer func() { + cerr := resp.Body.Close() + if cerr != nil && err == nil { + err = errors.Wrap(cerr, "closing response body") + } + }() + + // Read the response body. + result, err := ioutil.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading response") + } + + // Decode the translations. + transMap = make(map[string]uint64, len(keys)) + err = json.Unmarshal(result, &transMap) + if err != nil { + return nil, errors.Wrap(err, "json decoding") + } + + return transMap, nil +} + +func (c *InternalClient) CreateFieldKeysNode(ctx context.Context, uri *pilosa.URI, index string, field string, keys ...string) (transMap map[string]uint64, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.FindKeysFieldNode") + defer span.Finish() + + // Create HTTP request. + u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/index/%s/field/%s", index, field)) + q := u.Query() + q.Add("remote", "true") + u.RawQuery = q.Encode() + reqData, err := json.Marshal(keys) + if err != nil { + return nil, errors.Wrap(err, "marshalling request") + } + req, err := http.NewRequest("PUT", u.String(), bytes.NewReader(reqData)) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + // Apply headers. + req.Header.Set("Content-Length", strconv.Itoa(len(reqData))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + + // Send the request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, errors.Wrap(err, "executing request") + } + defer func() { + cerr := resp.Body.Close() + if cerr != nil && err == nil { + err = errors.Wrap(cerr, "closing response body") + } + }() + + // Read the response body. + result, err := ioutil.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading response") + } + + // Decode the translations. + transMap = make(map[string]uint64, len(keys)) + err = json.Unmarshal(result, &transMap) + if err != nil { + return nil, errors.Wrap(err, "json decoding") + } + + return transMap, nil +} + func (c *InternalClient) Transactions(ctx context.Context) (map[string]*pilosa.Transaction, error) { span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.Transactions") defer span.Finish() diff --git a/http/handler.go b/http/handler.go index d6162a5cb..ba0b9b13d 100644 --- a/http/handler.go +++ b/http/handler.go @@ -385,6 +385,11 @@ func newRouter(handler *Handler) *mux.Router { router.HandleFunc("/internal/nodes", handler.handleGetNodes).Methods("GET").Name("GetNodes") router.HandleFunc("/internal/shards/max", handler.handleGetShardsMax).Methods("GET").Name("GetShardsMax") // TODO: deprecate, but it's being used by the client + router.HandleFunc("/internal/translate/index/{index}/keys", handler.handleFindIndexKeys).Methods("POST").Name("FindIndexKeys") + router.HandleFunc("/internal/translate/index/{index}/keys", handler.handleCreateIndexKeys).Methods("PUT").Name("CreateIndexKeys") + router.HandleFunc("/internal/translate/field/{index}/{field}/keys", handler.handleFindFieldKeys).Methods("POST").Name("FindFieldKeys") + router.HandleFunc("/internal/translate/field/{index}/{field}/keys", handler.handleCreateFieldKeys).Methods("PUT").Name("CreateFieldKeys") + router.Use(handler.queryArgValidator) router.Use(handler.addQueryContext) router.Use(handler.extractTracing) @@ -2208,3 +2213,183 @@ func readBody(r *http.Request) ([]byte, error) { return buf.Bytes(), nil } + +func (h *Handler) handleFindIndexKeys(w http.ResponseWriter, r *http.Request) { + // Verify input and output types + if r.Header.Get("Content-Type") != "application/json" { + http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType) + return + } else if validHeaderAcceptJSON(r.Header) { + http.Error(w, "Not acceptable", http.StatusNotAcceptable) + return + } + + indexName, ok := mux.Vars(r)["index"] + if !ok { + http.Error(w, "index name is required", http.StatusBadRequest) + return + } + + bd, err := readBody(r) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + + var keys []string + err = json.Unmarshal(bd, &keys) + if err != nil { + http.Error(w, "failed to decode request", http.StatusBadRequest) + return + } + + translations, err := h.api.FindIndexKeys(r.Context(), indexName, keys...) + if err != nil { + http.Error(w, "translating keys", http.StatusBadRequest) + return + } + + err = json.NewEncoder(w).Encode(translations) + if err != nil { + http.Error(w, "encoding result", http.StatusBadRequest) + return + } +} + +func (h *Handler) handleFindFieldKeys(w http.ResponseWriter, r *http.Request) { + // Verify input and output types + if r.Header.Get("Content-Type") != "application/json" { + http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType) + return + } else if validHeaderAcceptJSON(r.Header) { + http.Error(w, "Not acceptable", http.StatusNotAcceptable) + return + } + + indexName, ok := mux.Vars(r)["index"] + if !ok { + http.Error(w, "index name is required", http.StatusBadRequest) + return + } + + fieldName, ok := mux.Vars(r)["field"] + if !ok { + http.Error(w, "field name is required", http.StatusBadRequest) + return + } + + bd, err := readBody(r) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + + var keys []string + err = json.Unmarshal(bd, &keys) + if err != nil { + http.Error(w, "failed to decode request", http.StatusBadRequest) + return + } + + translations, err := h.api.FindFieldKeys(r.Context(), indexName, fieldName, keys...) + if err != nil { + http.Error(w, "translating keys", http.StatusBadRequest) + return + } + + err = json.NewEncoder(w).Encode(translations) + if err != nil { + http.Error(w, "encoding result", http.StatusBadRequest) + return + } +} + +func (h *Handler) handleCreateIndexKeys(w http.ResponseWriter, r *http.Request) { + // Verify input and output types + if r.Header.Get("Content-Type") != "application/json" { + http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType) + return + } else if validHeaderAcceptJSON(r.Header) { + http.Error(w, "Not acceptable", http.StatusNotAcceptable) + return + } + + indexName, ok := mux.Vars(r)["index"] + if !ok { + http.Error(w, "index name is required", http.StatusBadRequest) + return + } + + bd, err := readBody(r) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + + var keys []string + err = json.Unmarshal(bd, &keys) + if err != nil { + http.Error(w, "failed to decode request", http.StatusBadRequest) + return + } + + translations, err := h.api.CreateIndexKeys(r.Context(), indexName, keys...) + if err != nil { + http.Error(w, "translating keys", http.StatusBadRequest) + return + } + + err = json.NewEncoder(w).Encode(translations) + if err != nil { + http.Error(w, "encoding result", http.StatusBadRequest) + return + } +} + +func (h *Handler) handleCreateFieldKeys(w http.ResponseWriter, r *http.Request) { + // Verify input and output types + if r.Header.Get("Content-Type") != "application/json" { + http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType) + return + } else if validHeaderAcceptJSON(r.Header) { + http.Error(w, "Not acceptable", http.StatusNotAcceptable) + return + } + + indexName, ok := mux.Vars(r)["index"] + if !ok { + http.Error(w, "index name is required", http.StatusBadRequest) + return + } + + fieldName, ok := mux.Vars(r)["field"] + if !ok { + http.Error(w, "field name is required", http.StatusBadRequest) + return + } + + bd, err := readBody(r) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + + var keys []string + err = json.Unmarshal(bd, &keys) + if err != nil { + http.Error(w, "failed to decode request", http.StatusBadRequest) + return + } + + translations, err := h.api.CreateFieldKeys(r.Context(), indexName, fieldName, keys...) + if err != nil { + http.Error(w, "translating keys", http.StatusBadRequest) + return + } + + err = json.NewEncoder(w).Encode(translations) + if err != nil { + http.Error(w, "encoding result", http.StatusBadRequest) + return + } +} diff --git a/mock/translator.go b/mock/translator.go index 8a28b504c..e7e88644f 100644 --- a/mock/translator.go +++ b/mock/translator.go @@ -21,8 +21,6 @@ import ( "github.com/pilosa/pilosa/v2" ) -var _ pilosa.TranslateStore = (*TranslateStore)(nil) - type TranslateStore struct { CloseFunc func() error MaxIDFunc func() (uint64, error) @@ -33,6 +31,8 @@ type TranslateStore struct { TranslateKeysFunc func(keys []string, writable bool) ([]uint64, error) TranslateIDFunc func(id uint64) (string, error) TranslateIDsFunc func(ids []uint64) ([]string, error) + FindKeysFunc func(keys ...string) (map[string]uint64, error) + CreateKeysFunc func(keys ...string) (map[string]uint64, error) ForceSetFunc func(id uint64, key string) error EntryReaderFunc func(ctx context.Context, offset uint64) (pilosa.TranslateEntryReader, error) } @@ -73,6 +73,14 @@ func (s *TranslateStore) TranslateIDs(ids []uint64) ([]string, error) { return s.TranslateIDsFunc(ids) } +func (s *TranslateStore) FindKeys(keys ...string) (map[string]uint64, error) { + return s.FindKeysFunc(keys...) +} + +func (s *TranslateStore) CreateKeys(keys ...string) (map[string]uint64, error) { + return s.CreateKeysFunc(keys...) +} + func (s *TranslateStore) ForceSet(id uint64, key string) error { return s.ForceSetFunc(id, key) } diff --git a/translate.go b/translate.go index 4efa26129..618f6e5d8 100644 --- a/translate.go +++ b/translate.go @@ -50,7 +50,7 @@ var ( // 0xc2, 0xa0, // NO-BREAK SPACE // 0x00, // } -type TranslateStore interface { +type TranslateStore interface { // TODO: refactor this interface; readonly should be part of the type and replication should be an impl detail io.Closer // Returns the maximum ID set on the store. @@ -71,6 +71,15 @@ type TranslateStore interface { TranslateKey(key string, writable bool) (uint64, error) TranslateKeys(key []string, writable bool) ([]uint64, error) + // FindKeys looks up the ID for each key. + // Keys are not created if they do not exist. + // Missing keys are not considered errors, so the length of the result may be less than that of the input. + FindKeys(keys ...string) (map[string]uint64, error) + + // CreateKeys maps all keys to IDs, creating the IDs if they do not exist. + // If the translator is read-only, this will return an error. + CreateKeys(keys ...string) (map[string]uint64, error) + // Converts an integer ID to its associated string key. TranslateID(id uint64) (string, error) TranslateIDs(id []uint64) ([]string, error) @@ -335,6 +344,57 @@ func (s *InMemTranslateStore) TranslateKeys(keys []string, writable bool) (_ []u return ids, nil } +// FindKeys looks up the ID for each key. +// Keys are not created if they do not exist. +// Missing keys are not considered errors, so the length of the result may be less than that of the input. +func (s *InMemTranslateStore) FindKeys(keys ...string) (map[string]uint64, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + result := make(map[string]uint64, len(keys)) + for _, key := range keys { + id, ok := s.idsByKey[key] + if !ok { + // The key does not exist. + continue + } + + result[key] = id + } + + return result, nil +} + +// CreateKeys maps all keys to IDs, creating the IDs if they do not exist. +// If the translator is read-only, this will return an error. +func (s *InMemTranslateStore) CreateKeys(keys ...string) (map[string]uint64, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.readOnly { + return nil, ErrTranslateStoreReadOnly + } + + result := make(map[string]uint64, len(keys)) + for _, key := range keys { + id, ok := s.idsByKey[key] + if !ok { + // The key does not exist. + // Generate a new id and update db. + if s.field == "" { + id = GenerateNextPartitionedID(s.index, s.maxID, s.partitionID, s.partitionN) + } else { + id = s.maxID + 1 + } + s.set(id, key) + } + + result[key] = id + } + + return result, nil +} + func (s *InMemTranslateStore) translateKey(key string, writable bool) (_ uint64, err error) { id := s.idsByKey[key] if id != 0 {