From 35672c53ce1ab09927d175f889c1e1c337bb626e Mon Sep 17 00:00:00 2001 From: Todd Gruben Date: Thu, 27 May 2021 09:10:12 -0500 Subject: [PATCH 1/2] added node partition endpoint;cluster aware key restore --- api.go | 17 +++++++ client.go | 4 ++ ctl/restore.go | 127 ++++++++++++++++++++++++++++++++++++++++++++---- http/client.go | 31 ++++++++++++ http/handler.go | 31 ++++++++++++ 5 files changed, 201 insertions(+), 9 deletions(-) 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..6272cf3fd 100644 --- a/ctl/restore.go +++ b/ctl/restore.go @@ -16,16 +16,19 @@ package ctl import ( "archive/tar" + "bytes" "compress/gzip" "context" "crypto/tls" "errors" "fmt" "io" + "io/ioutil" gohttp "net/http" "os" "strconv" "strings" + "sync" "github.com/pilosa/pilosa/v2" "github.com/pilosa/pilosa/v2/server" @@ -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,28 +148,95 @@ 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 err + return fmt.Errorf("cannot determine fragment nodes: %w", err) + } else if len(nodes) == 0 { + return fmt.Errorf("no nodes available") } + wg := new(sync.WaitGroup) + wg.Add(len(nodes)) + errors := make(chan error) + finished := make(chan bool) + + shardBytes, err := ioutil.ReadAll(tarReader) // this feels wrong but works for now + + for _, node := range nodes { + go func() { + defer wg.Done() + 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) + if err != nil { + errors <- err + } + }() + } + go func() { + wg.Wait() + close(finished) + }() + var problems error + select { + case <-finished: + break + case err := <-errors: + logger.Printf("error sending %v", err) + problems = err + } + if problems != nil { + return problems + } + case "translate": partitionID, err := strconv.Atoi(record[3]) logger.Printf("column keys %v (%v)", indexName, partitionID) if err != nil { return err } + partitionNodes, err := cmd.client.PartitionNodes(ctx, partitionID) + wg := new(sync.WaitGroup) + wg.Add(len(partitionNodes)) + errors := make(chan error) + finished := make(chan bool) - err = cmd.client.ImportIndexKeys(ctx, &primary.URI, indexName, partitionID, false, tarReader) + shardBytes, err := ioutil.ReadAll(tarReader) // this feels wrong but works for now if err != nil { return err } + for _, n := range partitionNodes { + node := n + go func() { + defer wg.Done() + rd := bytes.NewReader(shardBytes) + err = cmd.client.ImportIndexKeys(ctx, &node.URI, indexName, partitionID, false, rd) + if err != nil { + errors <- err + } + }() + } + go func() { + wg.Wait() + close(finished) + }() + var problems error + select { + case <-finished: + break + case err := <-errors: + logger.Printf("error sending %v", err) + problems = err + } + if problems != nil { + return problems + } + case "attributes": //skip case "fields": @@ -173,10 +244,48 @@ 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) + /* + err := cmd.client.ImportFieldKeys(ctx, &node.URI, indexName, fieldName, false, tarReader) + if err != nil { + return err + } + */ + //needs to go to all nodes + wg := new(sync.WaitGroup) + wg.Add(len(nodes)) + errors := make(chan error) + finished := make(chan bool) + + shardBytes, err := ioutil.ReadAll(tarReader) // this feels wrong but works for now if err != nil { return err } + for _, n := range nodes { + node := n + go func() { + defer wg.Done() + rd := bytes.NewReader(shardBytes) + err := cmd.client.ImportFieldKeys(ctx, &node.URI, indexName, fieldName, false, rd) + if err != nil { + errors <- err + } + }() + } + go func() { + wg.Wait() + close(finished) + }() + var problems error + select { + case <-finished: + break + case err := <-errors: + logger.Printf("error sending %v", err) + problems = err + } + if problems != nil { + return problems + } 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) { From c16a7a06a02d9b7faf553c402d5af8881e3d2de0 Mon Sep 17 00:00:00 2001 From: Todd Gruben Date: Fri, 28 May 2021 11:02:31 -0500 Subject: [PATCH 2/2] applied bens suggestions mainly using errgroup --- ctl/restore.go | 119 +++++++++++++------------------------------------ 1 file changed, 30 insertions(+), 89 deletions(-) diff --git a/ctl/restore.go b/ctl/restore.go index 6272cf3fd..6b149768c 100644 --- a/ctl/restore.go +++ b/ctl/restore.go @@ -28,11 +28,11 @@ import ( "os" "strconv" "strings" - "sync" "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 @@ -158,42 +158,26 @@ func (cmd *RestoreCommand) Run(ctx context.Context) (err error) { } else if len(nodes) == 0 { return fmt.Errorf("no nodes available") } - wg := new(sync.WaitGroup) - wg.Add(len(nodes)) - errors := make(chan error) - finished := make(chan bool) 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 { - go func() { - defer wg.Done() + 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) - if err != nil { - errors <- err - } - }() + return err + }) } - go func() { - wg.Wait() - close(finished) - }() - var problems error - select { - case <-finished: - break - case err := <-errors: - logger.Printf("error sending %v", err) - problems = err + if err := g.Wait(); err != nil { + return err } - if problems != nil { - return problems - } - case "translate": partitionID, err := strconv.Atoi(record[3]) logger.Printf("column keys %v (%v)", indexName, partitionID) @@ -201,40 +185,23 @@ func (cmd *RestoreCommand) Run(ctx context.Context) (err error) { return err } partitionNodes, err := cmd.client.PartitionNodes(ctx, partitionID) - wg := new(sync.WaitGroup) - wg.Add(len(partitionNodes)) - errors := make(chan error) - finished := make(chan bool) - + if err != nil { + return err + } shardBytes, err := ioutil.ReadAll(tarReader) // this feels wrong but works for now if err != nil { return err } - for _, n := range partitionNodes { - node := n - go func() { - defer wg.Done() + g, _ := errgroup.WithContext(ctx) + for _, node := range partitionNodes { + node := node + g.Go(func() error { rd := bytes.NewReader(shardBytes) - err = cmd.client.ImportIndexKeys(ctx, &node.URI, indexName, partitionID, false, rd) - if err != nil { - errors <- err - } - }() + return cmd.client.ImportIndexKeys(ctx, &node.URI, indexName, partitionID, false, rd) + }) } - go func() { - wg.Wait() - close(finished) - }() - var problems error - select { - case <-finished: - break - case err := <-errors: - logger.Printf("error sending %v", err) - problems = err - } - if problems != nil { - return problems + if err := g.Wait(); err != nil { + return err } case "attributes": @@ -244,47 +211,21 @@ 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, &node.URI, indexName, fieldName, false, tarReader) - if err != nil { - return err - } - */ //needs to go to all nodes - wg := new(sync.WaitGroup) - wg.Add(len(nodes)) - errors := make(chan error) - finished := make(chan bool) - shardBytes, err := ioutil.ReadAll(tarReader) // this feels wrong but works for now if err != nil { return err } - for _, n := range nodes { - node := n - go func() { - defer wg.Done() + g, _ := errgroup.WithContext(ctx) + for _, node := range nodes { + node := node + g.Go(func() error { rd := bytes.NewReader(shardBytes) - err := cmd.client.ImportFieldKeys(ctx, &node.URI, indexName, fieldName, false, rd) - if err != nil { - errors <- err - } - }() + return cmd.client.ImportFieldKeys(ctx, &node.URI, indexName, fieldName, false, rd) + }) } - go func() { - wg.Wait() - close(finished) - }() - var problems error - select { - case <-finished: - break - case err := <-errors: - logger.Printf("error sending %v", err) - problems = err - } - if problems != nil { - return problems + if err := g.Wait(); err != nil { + return err } case "attributes": //skip