diff --git a/api.go b/api.go index 715945339..5123f8afb 100644 --- a/api.go +++ b/api.go @@ -722,6 +722,21 @@ func (api *API) ShardNodes(ctx context.Context, indexName string, shard uint64) return snap.ShardNodes(indexName, shard), nil } +// PartitionNodes returns the node and all replicas which should contain a partition key data. +func (api *API) PartitionNodes(ctx context.Context, partitionID int) ([]*topology.Node, error) { + span, _ := tracing.StartSpanFromContext(ctx, "API.PartitionNodes") + defer span.Finish() + + if err := api.validate(apiPartitionNodes); err != nil { + return nil, errors.Wrap(err, "validating api method") + } + + // Create a snapshot of the cluster to use for node/partition calculations. + snap := topology.NewClusterSnapshot(api.cluster.noder, api.cluster.Hasher, api.cluster.ReplicaN) + + return snap.PartitionNodes(partitionID), nil +} + // FragmentBlockData is an endpoint for internal usage. It is not guaranteed to // return anything useful. Currently it returns protobuf encoded row and column // ids from a "block" which is a subdivision of a fragment. @@ -2372,6 +2387,7 @@ const ( apiIDReserve apiIDCommit apiIDReset + apiPartitionNodes ) var methodsCommon = map[apiMethod]struct{}{ @@ -2438,4 +2454,5 @@ var methodsNormal = map[apiMethod]struct{}{ apiIDReserve: {}, apiIDCommit: {}, apiIDReset: {}, + apiPartitionNodes: {}, } diff --git a/client.go b/client.go index 5143f2e6b..d06750aa0 100644 --- a/client.go +++ b/client.go @@ -59,6 +59,7 @@ type InternalClient interface { PostSchema(ctx context.Context, uri *pnet.URI, s *Schema, remote bool) error CreateIndex(ctx context.Context, index string, opt IndexOptions) error FragmentNodes(ctx context.Context, index string, shard uint64) ([]*topology.Node, error) + PartitionNodes(ctx context.Context, partitionID int) ([]*topology.Node, error) Nodes(ctx context.Context) ([]*topology.Node, error) Query(ctx context.Context, index string, queryRequest *QueryRequest) (*QueryResponse, error) Import(ctx context.Context, index, field string, shard uint64, bits []Bit, opts ...ImportOption) error @@ -189,6 +190,9 @@ func (n nopInternalClient) CreateIndex(ctx context.Context, index string, opt In func (n nopInternalClient) FragmentNodes(ctx context.Context, index string, shard uint64) ([]*topology.Node, error) { return nil, nil } +func (n nopInternalClient) PartitionNodes(ctx context.Context, partitionID int) ([]*topology.Node, error) { + return nil, nil +} func (n nopInternalClient) Nodes(ctx context.Context) ([]*topology.Node, error) { return nil, nil } diff --git a/ctl/restore.go b/ctl/restore.go index ae223fe1b..6b149768c 100644 --- a/ctl/restore.go +++ b/ctl/restore.go @@ -16,12 +16,14 @@ package ctl import ( "archive/tar" + "bytes" "compress/gzip" "context" "crypto/tls" "errors" "fmt" "io" + "io/ioutil" gohttp "net/http" "os" "strconv" @@ -30,6 +32,7 @@ import ( "github.com/pilosa/pilosa/v2" "github.com/pilosa/pilosa/v2/server" "github.com/pilosa/pilosa/v2/topology" + "golang.org/x/sync/errgroup" ) // RestoreCommand represents a command for restoring a backup to @@ -101,6 +104,7 @@ func (cmd *RestoreCommand) Run(ctx context.Context) (err error) { if err != nil { return err } + var primary *topology.Node for _, node := range nodes { if node.IsPrimary { @@ -109,10 +113,10 @@ func (cmd *RestoreCommand) Run(ctx context.Context) (err error) { } } - c := &gohttp.Client{} if primary == nil { return errors.New("no primary") } + c := &gohttp.Client{} for { header, err := tarReader.Next() if err == io.EOF { @@ -144,15 +148,34 @@ func (cmd *RestoreCommand) Run(ctx context.Context) (err error) { indexName := record[1] switch record[2] { case "shards": - shard, err := strconv.Atoi(record[3]) + shard, err := strconv.ParseUint(record[3], 10, 64) if err != nil { return err } - logger.Printf("shard %v %v", shard, indexName) - url := primary.URI.Path(fmt.Sprintf("/internal/restore/%v/%v", indexName, shard)) - //TODO (twg) cluster aware client - _, err = c.Post(url, "application/octet-stream", tarReader) + nodes, err := cmd.client.FragmentNodes(ctx, indexName, shard) if err != nil { + return fmt.Errorf("cannot determine fragment nodes: %w", err) + } else if len(nodes) == 0 { + return fmt.Errorf("no nodes available") + } + + shardBytes, err := ioutil.ReadAll(tarReader) // this feels wrong but works for now + if err != nil { + return err + } + g, _ := errgroup.WithContext(ctx) + for _, node := range nodes { + node := node + g.Go(func() error { + client := &gohttp.Client{} + rd := bytes.NewReader(shardBytes) + logger.Printf("shard %v %v", shard, indexName) + url := node.URI.Path(fmt.Sprintf("/internal/restore/%v/%v", indexName, shard)) + _, err = client.Post(url, "application/octet-stream", rd) + return err + }) + } + if err := g.Wait(); err != nil { return err } case "translate": @@ -161,11 +184,26 @@ func (cmd *RestoreCommand) Run(ctx context.Context) (err error) { if err != nil { return err } - - err = cmd.client.ImportIndexKeys(ctx, &primary.URI, indexName, partitionID, false, tarReader) + partitionNodes, err := cmd.client.PartitionNodes(ctx, partitionID) if err != nil { return err } + shardBytes, err := ioutil.ReadAll(tarReader) // this feels wrong but works for now + if err != nil { + return err + } + g, _ := errgroup.WithContext(ctx) + for _, node := range partitionNodes { + node := node + g.Go(func() error { + rd := bytes.NewReader(shardBytes) + return cmd.client.ImportIndexKeys(ctx, &node.URI, indexName, partitionID, false, rd) + }) + } + if err := g.Wait(); err != nil { + return err + } + case "attributes": //skip case "fields": @@ -173,10 +211,22 @@ func (cmd *RestoreCommand) Run(ctx context.Context) (err error) { switch action := record[4]; action { case "translate": logger.Printf("field keys %v %v", indexName, fieldName) - err := cmd.client.ImportFieldKeys(ctx, &primary.URI, indexName, fieldName, false, tarReader) + //needs to go to all nodes + shardBytes, err := ioutil.ReadAll(tarReader) // this feels wrong but works for now if err != nil { return err } + g, _ := errgroup.WithContext(ctx) + for _, node := range nodes { + node := node + g.Go(func() error { + rd := bytes.NewReader(shardBytes) + return cmd.client.ImportFieldKeys(ctx, &node.URI, indexName, fieldName, false, rd) + }) + } + if err := g.Wait(); err != nil { + return err + } case "attributes": //skip default: diff --git a/http/client.go b/http/client.go index 8fcff044d..29f943699 100644 --- a/http/client.go +++ b/http/client.go @@ -2157,3 +2157,34 @@ func (c *InternalClient) Status(ctx context.Context) (string, error) { } return rsp.State, nil } + +func (c *InternalClient) PartitionNodes(ctx context.Context, partitionID int) ([]*topology.Node, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.PartitionNodes") + defer span.Finish() + + // Execute request against the host. + u := uriPathToURL(c.defaultURI, "/internal/partition/nodes") + u.RawQuery = (url.Values{"partition": {strconv.FormatInt(int64(partitionID), 10)}}).Encode() + + // Build request. + req, err := http.NewRequest("GET", u.String(), nil) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + req.Header.Set("Accept", "application/json") + + // Execute request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + var a []*topology.Node + if err := json.NewDecoder(resp.Body).Decode(&a); err != nil { + return nil, fmt.Errorf("json decode: %s", err) + } + return a, nil +} diff --git a/http/handler.go b/http/handler.go index e5d23f7aa..444df39db 100644 --- a/http/handler.go +++ b/http/handler.go @@ -249,6 +249,7 @@ func (h *Handler) populateValidators() { h.validators["GetFragmentBlocks"] = queryValidationSpecRequired("index", "field", "view", "shard") h.validators["GetFragmentData"] = queryValidationSpecRequired("index", "field", "view", "shard") h.validators["GetFragmentNodes"] = queryValidationSpecRequired("shard", "index") + h.validators["GetPartitionNodes"] = queryValidationSpecRequired("partition") h.validators["GetNodes"] = queryValidationSpecRequired() h.validators["GetShardMax"] = queryValidationSpecRequired() h.validators["GetTransactionList"] = queryValidationSpecRequired() @@ -419,6 +420,7 @@ func newRouter(handler *Handler) http.Handler { router.HandleFunc("/internal/fragment/blocks", handler.handleGetFragmentBlocks).Methods("GET").Name("GetFragmentBlocks") router.HandleFunc("/internal/fragment/data", handler.handleGetFragmentData).Methods("GET").Name("GetFragmentData") router.HandleFunc("/internal/fragment/nodes", handler.handleGetFragmentNodes).Methods("GET").Name("GetFragmentNodes") + router.HandleFunc("/internal/partition/nodes", handler.handleGetPartitionNodes).Methods("GET").Name("GetPartitionNodes") router.HandleFunc("/internal/translate/data", handler.handleGetTranslateData).Methods("GET").Name("GetTranslateData") router.HandleFunc("/internal/translate/data", handler.handlePostTranslateData).Methods("POST").Name("PostTranslateData") router.HandleFunc("/internal/translate/keys", handler.handlePostTranslateKeys).Methods("POST").Name("PostTranslateKeys") @@ -1865,6 +1867,35 @@ func (h *Handler) handleGetFragmentNodes(w http.ResponseWriter, r *http.Request) } } +// handleGetPartitionNodes handles /internal/partition/nodes requests. +func (h *Handler) handleGetPartitionNodes(w http.ResponseWriter, r *http.Request) { + if !validHeaderAcceptJSON(r.Header) { + http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable) + return + } + q := r.URL.Query() + + // Read partition parameter. + partitionID, err := strconv.ParseInt(q.Get("partition"), 10, 64) + if err != nil { + http.Error(w, "shard should be an unsigned integer", http.StatusBadRequest) + return + } + + // Retrieve fragment owner nodes. + nodes, err := h.api.PartitionNodes(r.Context(), int(partitionID)) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + // Write to response. + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(nodes); err != nil { + h.logger.Errorf("json write error: %s", err) + } +} + // handleGetNodes handles /internal/nodes requests. func (h *Handler) handleGetNodes(w http.ResponseWriter, r *http.Request) { if !validHeaderAcceptJSON(r.Header) {