diff --git a/client.go b/client.go index 1f311a70a..b171b0dca 100644 --- a/client.go +++ b/client.go @@ -1,14 +1,19 @@ package pilosa import ( + "archive/tar" "bytes" "encoding/json" "errors" "fmt" + "io" "io/ioutil" + "log" + "math/rand" "net/http" "net/url" "strconv" + "time" "github.com/gogo/protobuf/proto" "github.com/umbel/pilosa/internal" @@ -38,6 +43,30 @@ func NewClient(host string) (*Client, error) { // Host returns the host the client was initialized with. func (c *Client) Host() string { return c.host } +// SliceN returns the number of slices on a server. +func (c *Client) SliceN() (uint64, error) { + // Execute request against the host. + u := url.URL{ + Scheme: "http", + Host: c.host, + Path: "/slices/max", + } + resp, err := c.HTTPClient.Get(u.String()) + if err != nil { + return 0, err + } + defer resp.Body.Close() + + var rsp sliceMaxResponse + if resp.StatusCode != http.StatusOK { + return 0, fmt.Errorf("http: status=%d", resp.StatusCode) + } else if err := json.NewDecoder(resp.Body).Decode(&rsp); err != nil { + return 0, fmt.Errorf("json decode: %s", err) + } + + return rsp.SliceMax, nil +} + // SliceNodes returns a list of nodes that own a slice. func (c *Client) SliceNodes(slice uint64) ([]*Node, error) { // Execute request against the host. @@ -140,6 +169,185 @@ func (c *Client) importNode(node *Node, buf []byte) error { return nil } +// BackupTo backs up an entire frame from a cluster to w. +func (c *Client) BackupTo(w io.Writer, db, frame string) error { + if db == "" { + return ErrDatabaseRequired + } else if frame == "" { + return ErrFrameRequired + } + + // Create tar writer around writer. + tw := tar.NewWriter(w) + + // Find the maximum number of slices. + sliceN, err := c.SliceN() + if err != nil { + return fmt.Errorf("slice n: %s", err) + } + + // Backup every slice to the tar file. + for i := uint64(0); i <= sliceN; i++ { + if err := c.backupSliceTo(tw, db, frame, i); err != nil { + return err + } + } + + // Close tar file. + if err := tw.Close(); err != nil { + return err + } + + return nil +} + +// backupSliceTo backs up a single slice to tw. +func (c *Client) backupSliceTo(tw *tar.Writer, db, frame string, slice uint64) error { + // Retrieve a list of nodes that own the slice. + nodes, err := c.SliceNodes(slice) + if err != nil { + return fmt.Errorf("slice nodes: %s", err) + } + + // Try to backup slice from each one until successful. + var data []byte + for _, i := range rand.Perm(len(nodes)) { + buf, err := c.backupSliceNode(db, frame, slice, nodes[i]) + if err == nil { + data = buf + break // backup successful + } else if err == ErrFragmentNotFound { + return nil // slice doesn't exist + } else if err != nil { + log.Println(err) + continue + } + } + + // Return error if unable to backup from any slice. + if data == nil { + return fmt.Errorf("unable to backup slice %d", slice) + } + + // Write slice file header. + if err := tw.WriteHeader(&tar.Header{ + Name: strconv.FormatUint(slice, 10), + Mode: 0666, + Size: int64(len(data)), + ModTime: time.Now(), + }); err != nil { + return err + } + + // Write buffer to file. + if _, err := tw.Write(data); err != nil { + return fmt.Errorf("write buffer: %s", err) + } + + return nil +} + +func (c *Client) backupSliceNode(db, frame string, slice uint64, node *Node) ([]byte, error) { + u := url.URL{ + Scheme: "http", + Host: node.Host, + Path: "/fragment/data", + RawQuery: url.Values{ + "db": {db}, + "frame": {frame}, + "slice": {strconv.FormatUint(slice, 10)}, + }.Encode(), + } + resp, err := c.HTTPClient.Get(u.String()) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + // Return error if status is not OK. + if resp.StatusCode == http.StatusNotFound { + return nil, ErrFragmentNotFound + } else if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("unexpected backup status code: host=%s, code=%d", node.Host, resp.StatusCode) + } + + return ioutil.ReadAll(resp.Body) +} + +// RestoreFrom restores a frame from a backup file to an entire cluster. +func (c *Client) RestoreFrom(r io.Reader, db, frame string) error { + if db == "" { + return ErrDatabaseRequired + } else if frame == "" { + return ErrFrameRequired + } + + // Create tar reader around input. + tr := tar.NewReader(r) + + // Process each file. + for { + hdr, err := tr.Next() + if err == io.EOF { + return nil + } else if err != nil { + return err + } + + // Parse slice from entry name. + slice, err := strconv.ParseUint(hdr.Name, 10, 64) + if err != nil { + return fmt.Errorf("invalid backup entry: %s", hdr.Name) + } + + // Read file into buffer. + var buf bytes.Buffer + if _, err := io.CopyN(&buf, tr, hdr.Size); err != nil { + return err + } + + // Restore file to all nodes that own it. + if err := c.restoreSliceFrom(buf.Bytes(), db, frame, slice); err != nil { + return err + } + } +} + +// restoreSliceFrom restores a single slice to all owning nodes. +func (c *Client) restoreSliceFrom(buf []byte, db, frame string, slice uint64) error { + // Retrieve a list of nodes that own the slice. + nodes, err := c.SliceNodes(slice) + if err != nil { + return fmt.Errorf("slice nodes: %s", err) + } + + // Restore slice to each owner. + for _, node := range nodes { + u := url.URL{ + Scheme: "http", + Host: node.Host, + Path: "/fragment/data", + RawQuery: url.Values{ + "db": {db}, + "frame": {frame}, + "slice": {strconv.FormatUint(slice, 10)}, + }.Encode(), + } + resp, err := c.HTTPClient.Post(u.String(), "application/octet-stream", bytes.NewReader(buf)) + if err != nil { + return err + } + resp.Body.Close() + + // Return error if response not OK. + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("unexpected status code: host=%s, code=%d", node.Host, resp.StatusCode) + } + } + + return nil +} + // Bit represents the location of a single bit. type Bit struct { BitmapID uint64 diff --git a/client_test.go b/client_test.go index 42ecb7143..3a4f24ab1 100644 --- a/client_test.go +++ b/client_test.go @@ -1,6 +1,7 @@ package pilosa_test import ( + "bytes" "reflect" "testing" @@ -39,6 +40,51 @@ func TestClient_Import(t *testing.T) { } } +// Ensure client backup and restore a frame. +func TestClient_BackupRestore(t *testing.T) { + idx := MustOpenIndex() + defer idx.Close() + + idx.MustCreateFragmentIfNotExists("d", "f", 0).MustSetBits(100, 1, 2, 3, SliceWidth-1) + idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBits(100, SliceWidth, SliceWidth+2) + idx.MustCreateFragmentIfNotExists("d", "f", 5).MustSetBits(100, (5*SliceWidth)+1) + idx.MustCreateFragmentIfNotExists("d", "f", 0).MustSetBits(200, 20000) + + s := NewServer() + defer s.Close() + s.Handler.Host = s.Host() + s.Handler.Cluster = NewCluster(1) + s.Handler.Cluster.Nodes[0].Host = s.Host() + s.Handler.Index = idx.Index + + c := MustNewClient(s.Host()) + + // Backup from frame. + var buf bytes.Buffer + if err := c.BackupTo(&buf, "d", "f"); err != nil { + t.Fatal(err) + } + + // Restore to a different frame. + if err := c.RestoreFrom(&buf, "x", "y"); err != nil { + t.Fatal(err) + } + + // Verify data. + if a := idx.Fragment("x", "y", 0).Bitmap(100).Bits(); !reflect.DeepEqual(a, []uint64{1, 2, 3, SliceWidth - 1}) { + t.Fatalf("unexpected bits(0): %+v", a) + } + if a := idx.Fragment("x", "y", 1).Bitmap(100).Bits(); !reflect.DeepEqual(a, []uint64{SliceWidth, SliceWidth + 2}) { + t.Fatalf("unexpected bits(0): %+v", a) + } + if a := idx.Fragment("x", "y", 5).Bitmap(100).Bits(); !reflect.DeepEqual(a, []uint64{(5 * SliceWidth) + 1}) { + t.Fatalf("unexpected bits(0): %+v", a) + } + if a := idx.Fragment("x", "y", 0).Bitmap(200).Bits(); !reflect.DeepEqual(a, []uint64{20000}) { + t.Fatalf("unexpected bits: %+v", a) + } +} + // Client represents a test wrapper for pilosa.Client. type Client struct { *pilosa.Client diff --git a/cmd/pilosactl/main.go b/cmd/pilosactl/main.go index 38f1d0717..549a38c58 100644 --- a/cmd/pilosactl/main.go +++ b/cmd/pilosactl/main.go @@ -78,6 +78,10 @@ func (m *Main) Run() error { cmd = NewConfigCommand(m.Stdin, m.Stdout, m.Stderr) case "import": cmd = NewImportCommand(m.Stdin, m.Stdout, m.Stderr) + case "backup": + cmd = NewBackupCommand(m.Stdin, m.Stdout, m.Stderr) + case "restore": + cmd = NewRestoreCommand(m.Stdin, m.Stdout, m.Stderr) default: return ErrUnknownCommand } @@ -324,3 +328,175 @@ func (cmd *ImportCommand) parsePath(path string) ([]pilosa.Bit, error) { return a, nil } + +// BackupCommand represents a command for backing up a frame. +type BackupCommand struct { + // Destination host and port. + Host string + + // Name of the database & frame to backup. + Database string + Frame string + + // Output file to write to. + Path string + + // Standard input/output + Stdin io.Reader + Stdout io.Writer + Stderr io.Writer +} + +// NewBackupCommand returns a new instance of BackupCommand. +func NewBackupCommand(stdin io.Reader, stdout, stderr io.Writer) *BackupCommand { + return &BackupCommand{ + Stdin: stdin, + Stdout: stdout, + Stderr: stderr, + } +} + +// ParseFlags parses command line flags from args. +func (cmd *BackupCommand) ParseFlags(args []string) error { + fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError) + fs.SetOutput(cmd.Stderr) + fs.StringVar(&cmd.Host, "host", "localhost:15000", "host:port") + fs.StringVar(&cmd.Database, "d", "", "database") + fs.StringVar(&cmd.Frame, "f", "", "frame") + fs.StringVar(&cmd.Path, "o", "", "output file") + if err := fs.Parse(args); err != nil { + return err + } + + return nil +} + +// Usage returns the usage message to be printed. +func (cmd *BackupCommand) Usage() string { + return strings.TrimSpace(` +usage: pilosactl backup -host HOST -d database -f frame -o PATH + +Backs up the database and frame from across the cluster into a single file. +`) +} + +// Run executes the main program execution. +func (cmd *BackupCommand) Run() error { + // Validate arguments. + if cmd.Path == "" { + return errors.New("output file required") + } + + // Create a client to the server. + client, err := pilosa.NewClient(cmd.Host) + if err != nil { + return err + } + + // Open output file. + f, err := os.Create(cmd.Path) + if err != nil { + return err + } + defer f.Close() + + // Begin streaming backup. + if err := client.BackupTo(f, cmd.Database, cmd.Frame); err != nil { + return err + } + + // Sync & close file to ensure durability. + if err := f.Sync(); err != nil { + return err + } else if err = f.Close(); err != nil { + return err + } + + return nil +} + +// RestoreCommand represents a command for restoring a frame from a backup. +type RestoreCommand struct { + // Destination host and port. + Host string + + // Name of the database & frame to backup. + Database string + Frame string + + // Import file to read from. + Path string + + // Standard input/output + Stdin io.Reader + Stdout io.Writer + Stderr io.Writer +} + +// NewRestoreCommand returns a new instance of RestoreCommand. +func NewRestoreCommand(stdin io.Reader, stdout, stderr io.Writer) *RestoreCommand { + return &RestoreCommand{ + Stdin: stdin, + Stdout: stdout, + Stderr: stderr, + } +} + +// ParseFlags parses command line flags from args. +func (cmd *RestoreCommand) ParseFlags(args []string) error { + fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError) + fs.SetOutput(cmd.Stderr) + fs.StringVar(&cmd.Host, "host", "localhost:15000", "host:port") + fs.StringVar(&cmd.Database, "d", "", "database") + fs.StringVar(&cmd.Frame, "f", "", "frame") + if err := fs.Parse(args); err != nil { + return err + } + + // Read input path from the args. + if fs.NArg() == 0 { + return errors.New("path required") + } else if fs.NArg() > 1 { + return errors.New("too many paths specified") + } + cmd.Path = fs.Arg(0) + + return nil +} + +// Usage returns the usage message to be printed. +func (cmd *RestoreCommand) Usage() string { + return strings.TrimSpace(` +usage: pilosactl restore -host HOST -d database -f frame PATH + +Restores a frame to the cluster from a backup file. +`) +} + +// Run executes the main program execution. +func (cmd *RestoreCommand) Run() error { + // Validate arguments. + if cmd.Path == "" { + return errors.New("backup file required") + } + + // Create a client to the server. + client, err := pilosa.NewClient(cmd.Host) + if err != nil { + return err + } + + // Open backup file. + f, err := os.Open(cmd.Path) + if err != nil { + return err + } + defer f.Close() + + // Restore backup file to the cluster. + if err := client.RestoreFrom(f, cmd.Database, cmd.Frame); err != nil { + return err + } + + return nil +} diff --git a/executor_test.go b/executor_test.go index 676cc20f1..2c1be9358 100644 --- a/executor_test.go +++ b/executor_test.go @@ -128,7 +128,7 @@ func TestExecutor_Execute_SetBit(t *testing.T) { t.Fatalf("unexpected bitmap count: %d", n) } - if res, err := e.Execute("d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil); err != nil { + if res, err := e.Execute("d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil, nil); err != nil { t.Fatal(err) } else { if !res.(bool) { @@ -139,7 +139,7 @@ func TestExecutor_Execute_SetBit(t *testing.T) { if n := f.Bitmap(11).Count(); n != 1 { t.Fatalf("unexpected bitmap count: %d", n) } - if res, err := e.Execute("d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil); err != nil { + if res, err := e.Execute("d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil, nil); err != nil { t.Fatal(err) } else { if res.(bool) { diff --git a/handler.go b/handler.go index 188473640..bb59d40d1 100644 --- a/handler.go +++ b/handler.go @@ -152,7 +152,6 @@ func (h *Handler) handlePostQuery(w http.ResponseWriter, r *http.Request) { } func (h *Handler) handleGetSliceMax(w http.ResponseWriter, r *http.Request) error { - sm := h.Index.SliceN() if strings.Contains(r.Header.Get("Accept"), "application/x-protobuf") { pb := &internal.SliceMaxResponse{ @@ -165,8 +164,11 @@ func (h *Handler) handleGetSliceMax(w http.ResponseWriter, r *http.Request) erro } return nil } - resp := map[string]uint64{"SliceMax": sm} - return json.NewEncoder(w).Encode(resp) + return json.NewEncoder(w).Encode(sliceMaxResponse{SliceMax: sm}) +} + +type sliceMaxResponse struct { + SliceMax uint64 `json:"SliceMax"` } // readProfiles returns a list of profile objects by id. diff --git a/pilosa.go b/pilosa.go index f788a2037..20ab8ca8a 100644 --- a/pilosa.go +++ b/pilosa.go @@ -18,6 +18,9 @@ var ( // ErrFrameRequired is returned when no frame is specified. ErrFrameRequired = errors.New("frame required") + + // ErrFragmentNotFound is returned when a fragment does not exist. + ErrFragmentNotFound = errors.New("fragment not found") ) // Version represents the current running version of Pilosa.