From 320e9811d8dc53fa1e4007ef187870b3e188b4e3 Mon Sep 17 00:00:00 2001 From: Ben Johnson Date: Tue, 29 Dec 2015 16:11:34 -0700 Subject: [PATCH] bulk importer This commit adds the pilosactl binary with a single `import` command. The import allows users to bulk import data from a CSV file. --- client.go | 134 +++++++++++++++++++++++++ client_test.go | 54 ++++++++++ cluster.go | 5 + cmd/pilosa/main.go | 6 ++ cmd/pilosa/main_test.go | 6 +- cmd/pilosactl/main.go | 217 ++++++++++++++++++++++++++++++++++++++-- executor.go | 20 ++-- executor_test.go | 52 +++++----- fragment.go | 209 ++++++++++++++++++++++++++++---------- fragment_test.go | 29 ++++++ handler.go | 98 ++++++++++++++++-- handler_test.go | 2 +- index.go | 86 +++++++++++++++- index_test.go | 6 +- internal/internal.pb.go | 91 +++++++++++++++++ internal/internal.proto | 17 ++++ pilosa.go | 15 +++ 17 files changed, 931 insertions(+), 116 deletions(-) create mode 100644 client.go create mode 100644 client_test.go diff --git a/client.go b/client.go new file mode 100644 index 000000000..d6676e968 --- /dev/null +++ b/client.go @@ -0,0 +1,134 @@ +package pilosa + +import ( + "bytes" + "errors" + "fmt" + "io/ioutil" + "net/http" + "net/url" + "strconv" + + "github.com/gogo/protobuf/proto" + "github.com/umbel/pilosa/internal" +) + +// Client represents a client to the Pilosa cluster. +type Client struct { + host string + + // The client to use for HTTP communication. + // Defaults to the http.DefaultClient. + HTTPClient *http.Client +} + +// NewClient returns a new instance of Client to connect to host. +func NewClient(host string) (*Client, error) { + if host == "" { + return nil, ErrHostRequired + } + + return &Client{ + host: host, + HTTPClient: http.DefaultClient, + }, nil +} + +// Host returns the host the client was initialized with. +func (c *Client) Host() string { return c.host } + +// Import bulk imports bits for a single slice to a host. +func (c *Client) Import(db, frame string, slice uint64, bits []Bit) error { + if db == "" { + return ErrDatabaseRequired + } else if frame == "" { + return ErrFrameRequired + } + + // Separate bitmap and profile IDs to reduce allocations. + bitmapIDs := Bits(bits).BitmapIDs() + profileIDs := Bits(bits).ProfileIDs() + + // Marshal bits to protobufs. + buf, err := proto.Marshal(&internal.ImportRequest{ + DB: proto.String(db), + Frame: proto.String(frame), + Slice: proto.Uint64(slice), + BitmapIDs: bitmapIDs, + ProfileIDs: profileIDs, + }) + if err != nil { + return fmt.Errorf("marshal import request: %s", err) + } + + // Create URL & HTTP request. + u := url.URL{Scheme: "http", Host: c.host, Path: "/import"} + req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf)) + if err != nil { + return err + } + req.Header.Set("Content-Length", strconv.Itoa(len(buf))) + req.Header.Set("Content-Type", "application/x-protobuf") + req.Header.Set("Accept", "application/x-protobuf") + + // Execute request against the host. + resp, err := c.HTTPClient.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + // Read body and unmarshal response. + body, err := ioutil.ReadAll(resp.Body) + if err != nil { + return err + } else if resp.StatusCode != http.StatusOK { + return errors.New(string(body)) + } + + var isresp internal.ImportResponse + if err := proto.Unmarshal(body, &isresp); err != nil { + return fmt.Errorf("unmarshal import response: %s", err) + } else if s := isresp.GetErr(); s != "" { + return errors.New(s) + } + + return nil +} + +// Bit represents the location of a single bit. +type Bit struct { + BitmapID uint64 + ProfileID uint64 +} + +// Bits represents a slice of bits. +type Bits []Bit + +// BitmapIDs returns a slice of all the bitmap IDs. +func (a Bits) BitmapIDs() []uint64 { + other := make([]uint64, len(a)) + for i := range a { + other[i] = a[i].BitmapID + } + return other +} + +// ProfileIDs returns a slice of all the profile IDs. +func (a Bits) ProfileIDs() []uint64 { + other := make([]uint64, len(a)) + for i := range a { + other[i] = a[i].ProfileID + } + return other +} + +// GroupBySlice returns a map of bits by slice. +func (a Bits) GroupBySlice() map[uint64][]Bit { + m := make(map[uint64][]Bit) + for _, bit := range a { + slice := bit.ProfileID / SliceWidth + m[slice] = append(m[slice], bit) + } + return m +} diff --git a/client_test.go b/client_test.go new file mode 100644 index 000000000..42ecb7143 --- /dev/null +++ b/client_test.go @@ -0,0 +1,54 @@ +package pilosa_test + +import ( + "reflect" + "testing" + + "github.com/umbel/pilosa" +) + +// Ensure client can bulk import data. +func TestClient_Import(t *testing.T) { + idx := MustOpenIndex() + defer idx.Close() + + 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 + + // Send import request. + c := MustNewClient(s.Host()) + if err := c.Import("d", "f", 0, []pilosa.Bit{ + {BitmapID: 0, ProfileID: 1}, + {BitmapID: 0, ProfileID: 5}, + {BitmapID: 200, ProfileID: 6}, + }); err != nil { + t.Fatal(err) + } + + // Verify data. + f := idx.MustCreateFragmentIfNotExists("d", "f", 0) + if a := f.Bitmap(0).Bits(); !reflect.DeepEqual(a, []uint64{1, 5}) { + t.Fatalf("unexpected bits: %+v", a) + } + if a := f.Bitmap(200).Bits(); !reflect.DeepEqual(a, []uint64{6}) { + t.Fatalf("unexpected bits: %+v", a) + } +} + +// Client represents a test wrapper for pilosa.Client. +type Client struct { + *pilosa.Client +} + +// MustNewClient returns a new instance of Client. Panic on error. +func MustNewClient(host string) *Client { + c, err := pilosa.NewClient(host) + if err != nil { + panic(err) + } + return &Client{Client: c} +} diff --git a/cluster.go b/cluster.go index efdadef98..cbb41fd0b 100644 --- a/cluster.go +++ b/cluster.go @@ -79,6 +79,11 @@ func (c *Cluster) SliceNodes(slice uint64) []*Node { return c.PartitionNodes(c.Partition(slice)) } +// OwnsSlice returns true if a host owns slice. +func (c *Cluster) OwnsSlice(host string, slice uint64) bool { + return Nodes(c.SliceNodes(slice)).ContainsHost(host) +} + // PartitionNodes returns a list of nodes that own a partition. func (c *Cluster) PartitionNodes(partitionID int) []*Node { // Default replica count to between one and the number of nodes. diff --git a/cmd/pilosa/main.go b/cmd/pilosa/main.go index 83f372e54..5061b2d0a 100644 --- a/cmd/pilosa/main.go +++ b/cmd/pilosa/main.go @@ -141,6 +141,9 @@ func (m *Main) Run(args ...string) error { // Create index to store fragments. fmt.Fprintf(m.Stderr, "Using data from: %s\n", m.Config.DataDir) m.index = pilosa.NewIndex(m.Config.DataDir) + if err := m.index.Open(); err != nil { + return err + } // Create executor for executing queries. e := pilosa.NewExecutor(m.index) @@ -149,6 +152,9 @@ func (m *Main) Run(args ...string) error { // Initialize HTTP handler. h := pilosa.NewHandler() + h.Index = m.index + h.Host = hostname + h.Cluster = cluster h.Executor = e h.LogOutput = m.Stderr diff --git a/cmd/pilosa/main_test.go b/cmd/pilosa/main_test.go index 93b0b3399..fe38a231f 100644 --- a/cmd/pilosa/main_test.go +++ b/cmd/pilosa/main_test.go @@ -21,6 +21,10 @@ import ( // Ensure program can process queries and maintain consistency. func TestMain_Set_Quick(t *testing.T) { + if testing.Short() { + t.Skip("short") + } + if err := quick.Check(func(cmds []SetCommand) bool { m := MustRunMain() defer m.Close() @@ -57,7 +61,7 @@ func TestMain_Set_Quick(t *testing.T) { if res, err := m.Query("d", fmt.Sprintf(`get(id=%d, frame=%q)`, id, frame)); err != nil { t.Fatal(err) } else if res != exp { - t.Fatalf("unexpected result:\n\ngot=%s\n\nexp=%s\n\n", res, exp) + t.Fatalf("unexpected result (reopen):\n\ngot=%s\n\nexp=%s\n\n", res, exp) } } } diff --git a/cmd/pilosactl/main.go b/cmd/pilosactl/main.go index 56631a35d..efacdd6d1 100644 --- a/cmd/pilosactl/main.go +++ b/cmd/pilosactl/main.go @@ -1,11 +1,32 @@ package main import ( + "encoding/csv" "errors" "flag" "fmt" "io" + "log" "os" + "strconv" + "strings" + + "github.com/umbel/pilosa" +) + +var ( + // ErrUsage is returned when usage should be displayed for the program. + ErrUsage = errors.New("usage") + + // ErrUnknownCommand is returned when specifying an unknown command. + ErrUnknownCommand = errors.New("unknown command") + + // ErrQuit is returned when the program should simply quit. + // This is used when the error message has already been printed. + ErrQuit = errors.New("quit") + + // ErrPathRequired is returned when executing a command without a required path. + ErrPathRequired = errors.New("path required") ) func main() { @@ -18,7 +39,9 @@ func main() { } // Execute the program. - if err := m.Run(); err != nil { + if err := m.Run(); err == ErrQuit { + os.Exit(1) + } else if err != nil { fmt.Fprintln(m.Stderr, err) os.Exit(1) } @@ -26,7 +49,9 @@ func main() { // Main represents the main program execution. type Main struct { + // Command name and arguments passed into the CLI. Command string + Args []string // Standard input/output Stdin io.Reader @@ -45,21 +70,201 @@ func NewMain() *Main { // Run executes the main program execution. func (m *Main) Run() error { - panic("FIXME: implement commands") + var cmd Command + switch m.Command { + case "", "help", "-h": + return ErrUsage + case "import": + cmd = NewImportCommand(m.Stdin, m.Stdout, m.Stderr) + default: + return ErrUnknownCommand + } + + // Parse command's flags. + if err := cmd.ParseFlags(m.Args); err == ErrUsage { + fmt.Fprintln(m.Stderr, cmd.Usage()) + return ErrQuit + } else if err != nil { + return err + } + + // Execute the command. + if err := cmd.Run(); err != nil { + return err + } + return nil } // ParseFlags parses command line flags from args. func (m *Main) ParseFlags(args []string) error { + if len(args) == 0 { + return nil + } + + m.Command = args[0] + m.Args = args[1:] + return nil +} + +// Command represents an executable subcommand. +type Command interface { + Usage() string + ParseFlags(args []string) error + Run() error +} + +// ImportCommand represents a command for bulk importing data. +type ImportCommand struct { + // Destination host and port. + Host string + + // Name of the database & frame to import into. + Database string + Frame string + + // Filenames to import from. + Paths []string + + // Standard input/output + Stdin io.Reader + Stdout io.Writer + Stderr io.Writer +} + +// NewImportCommand returns a new instance of ImportCommand. +func NewImportCommand(stdin io.Reader, stdout, stderr io.Writer) *ImportCommand { + return &ImportCommand{ + Stdin: stdin, + Stdout: stdout, + Stderr: stderr, + } +} + +// ParseFlags parses command line flags from args. +func (cmd *ImportCommand) ParseFlags(args []string) error { fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError) - fs.SetOutput(m.Stderr) + 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 } - // Require command. - if fs.NArg() == 0 { - return errors.New("command required") + // Extract the import paths. + cmd.Paths = fs.Args() + + return nil +} + +// Usage returns the usage message to be printed. +func (cmd *ImportCommand) Usage() string { + return strings.TrimSpace(` +usage: pilosactl import -host HOST -d database -f frame paths + +Bulk imports one or more CSV files to a host's database and frame. The bits +of the CSV file are grouped by slice for the most efficient import. + +The format of the CSV file is: + + BITMAPID,PROFILEID + +The file should contain no headers. +`) +} + +// Run executes the main program execution. +func (cmd *ImportCommand) Run() error { + logger := log.New(cmd.Stderr, "", log.LstdFlags) + + // Validate arguments. + // Database and frame are validated early before the files are parsed. + if cmd.Database == "" { + return pilosa.ErrDatabaseRequired + } else if cmd.Frame == "" { + return pilosa.ErrFrameRequired + } else if len(cmd.Paths) == 0 { + return ErrPathRequired + } + + // Create a client to the server. + client, err := pilosa.NewClient(cmd.Host) + if err != nil { + return err + } + + // Import each path and import by slice. + for _, path := range cmd.Paths { + // Parse path into bits. + logger.Printf("parsing: %s", path) + bits, err := cmd.parsePath(path) + if err != nil { + return err + } + + // Group bits by slice. + logger.Printf("grouping %d bits", len(bits)) + bitsBySlice := pilosa.Bits(bits).GroupBySlice() + logger.Printf("grouped into %d slices", len(bitsBySlice)) + + // Parse path into bits. + for slice, bits := range bitsBySlice { + logger.Printf("importing slice: %d, n=%d", slice, len(bits)) + if err := client.Import(cmd.Database, cmd.Frame, slice, bits); err != nil { + return err + } + } } return nil } + +// parsePath parses a path into bits. +func (cmd *ImportCommand) parsePath(path string) ([]pilosa.Bit, error) { + var a []pilosa.Bit + + // Open file for reading. + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + // Read rows as bits. + r := csv.NewReader(f) + rnum := 0 + for { + rnum++ + + // Read CSV row. + record, err := r.Read() + if err == io.EOF { + break + } else if err != nil { + return nil, err + } + + // Ignore blank rows. + if record[0] == "" { + continue + } else if len(record) != 2 { + return nil, fmt.Errorf("bad column count on row %d: col=%d", rnum, len(record)) + } + + // Parse bitmap id. + bitmapID, err := strconv.ParseUint(record[0], 10, 64) + if err != nil { + return nil, fmt.Errorf("invalid bitmap id on row %d: %q", rnum, record[0]) + } + + // Parse bitmap id. + profileID, err := strconv.ParseUint(record[1], 10, 64) + if err != nil { + return nil, fmt.Errorf("invalid profile id on row %d: %q", rnum, record[1]) + } + + a = append(a, pilosa.Bit{BitmapID: bitmapID, ProfileID: profileID}) + } + + return a, nil +} diff --git a/executor.go b/executor.go index 4ab485f0b..d1a84f783 100644 --- a/executor.go +++ b/executor.go @@ -16,17 +16,12 @@ import ( // DefaultFrame is the frame used if one is not specified. const DefaultFrame = "general" -// ErrDatabaseRequired is returned when no database is specified. -var ErrDatabaseRequired = errors.New("database required") - // Executor recursively executes calls in a PQL query across all slices. type Executor struct { index *Index - // Local hostname. - Host string - - // Cluster configuration + // Local hostname & cluster configuration. + Host string Cluster *Cluster // Client used for remote HTTP requests. @@ -60,7 +55,8 @@ func (e *Executor) Execute(db string, q *pql.Query, slices []uint64) (interface{ // If slices aren't specified, then include all of them. if len(slices) == 0 { // Round up the number of slices. - sliceN := (e.index.SliceN() % uint64(len(e.Cluster.Nodes))) + uint64(len(e.Cluster.Nodes)) + sliceN := e.index.SliceN() + sliceN += (sliceN % uint64(len(e.Cluster.Nodes))) + uint64(len(e.Cluster.Nodes)) // Generate a slices of all slices. slices = make([]uint64, sliceN+1) @@ -160,9 +156,9 @@ func (e *Executor) executeGetSlice(db string, c *pql.Get, slice uint64) (*Bitmap frame = DefaultFrame } - f, err := e.Index().Fragment(db, frame, slice) - if err != nil { - return nil, fmt.Errorf("fragment: %s", err) + f := e.Index().Fragment(db, frame, slice) + if f == nil { + return NewBitmap(), nil } return f.Bitmap(c.ID), nil } @@ -243,7 +239,7 @@ func (e *Executor) executeSet(db string, c *pql.Set) error { for _, node := range e.Cluster.SliceNodes(slice) { // Update locally if host matches. if node.Host == e.Host { - f, err := e.Index().Fragment(db, c.Frame, slice) + f, err := e.Index().CreateFragmentIfNotExists(db, c.Frame, slice) if err != nil { return fmt.Errorf("fragment: %s", err) } diff --git a/executor_test.go b/executor_test.go index 817434795..530f4c162 100644 --- a/executor_test.go +++ b/executor_test.go @@ -14,8 +14,8 @@ import ( func TestExecutor_Execute_Get(t *testing.T) { idx := MustOpenIndex() defer idx.Close() - idx.MustFragment("d", "f", 0).MustSetBit(10, 3) - idx.MustFragment("d", "f", 1).MustSetBit(10, SliceWidth+1) + idx.MustCreateFragmentIfNotExists("d", "f", 0).MustSetBit(10, 3) + idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBit(10, SliceWidth+1) e := NewExecutor(idx.Index, NewCluster(1)) if res, err := e.Execute("d", MustParse(`get(id=10, frame=f)`), nil); err != nil { @@ -33,10 +33,10 @@ func TestExecutor_Execute_Get(t *testing.T) { func TestExecutor_Execute_Difference(t *testing.T) { idx := MustOpenIndex() defer idx.Close() - idx.MustFragment("d", "general", 0).MustSetBit(10, 1) - idx.MustFragment("d", "general", 0).MustSetBit(10, 2) - idx.MustFragment("d", "general", 0).MustSetBit(10, 3) - idx.MustFragment("d", "general", 0).MustSetBit(11, 2) + idx.MustCreateFragmentIfNotExists("d", "general", 0).MustSetBit(10, 1) + idx.MustCreateFragmentIfNotExists("d", "general", 0).MustSetBit(10, 2) + idx.MustCreateFragmentIfNotExists("d", "general", 0).MustSetBit(10, 3) + idx.MustCreateFragmentIfNotExists("d", "general", 0).MustSetBit(11, 2) e := NewExecutor(idx.Index, NewCluster(1)) if res, err := e.Execute("d", MustParse(`difference(get(id=10), get(id=11))`), nil); err != nil { @@ -52,13 +52,13 @@ func TestExecutor_Execute_Difference(t *testing.T) { func TestExecutor_Execute_Intersect(t *testing.T) { idx := MustOpenIndex() defer idx.Close() - idx.MustFragment("d", "general", 0).MustSetBit(10, 1) - idx.MustFragment("d", "general", 1).MustSetBit(10, SliceWidth+1) - idx.MustFragment("d", "general", 1).MustSetBit(10, SliceWidth+2) + idx.MustCreateFragmentIfNotExists("d", "general", 0).MustSetBit(10, 1) + idx.MustCreateFragmentIfNotExists("d", "general", 1).MustSetBit(10, SliceWidth+1) + idx.MustCreateFragmentIfNotExists("d", "general", 1).MustSetBit(10, SliceWidth+2) - idx.MustFragment("d", "general", 0).MustSetBit(11, 1) - idx.MustFragment("d", "general", 0).MustSetBit(11, 2) - idx.MustFragment("d", "general", 1).MustSetBit(11, SliceWidth+2) + idx.MustCreateFragmentIfNotExists("d", "general", 0).MustSetBit(11, 1) + idx.MustCreateFragmentIfNotExists("d", "general", 0).MustSetBit(11, 2) + idx.MustCreateFragmentIfNotExists("d", "general", 1).MustSetBit(11, SliceWidth+2) e := NewExecutor(idx.Index, NewCluster(1)) if res, err := e.Execute("d", MustParse(`intersect(get(id=10), get(id=11))`), nil); err != nil { @@ -76,12 +76,12 @@ func TestExecutor_Execute_Intersect(t *testing.T) { func TestExecutor_Execute_Union(t *testing.T) { idx := MustOpenIndex() defer idx.Close() - idx.MustFragment("d", "general", 0).MustSetBit(10, 0) - idx.MustFragment("d", "general", 1).MustSetBit(10, SliceWidth+1) - idx.MustFragment("d", "general", 1).MustSetBit(10, SliceWidth+2) + idx.MustCreateFragmentIfNotExists("d", "general", 0).MustSetBit(10, 0) + idx.MustCreateFragmentIfNotExists("d", "general", 1).MustSetBit(10, SliceWidth+1) + idx.MustCreateFragmentIfNotExists("d", "general", 1).MustSetBit(10, SliceWidth+2) - idx.MustFragment("d", "general", 0).MustSetBit(11, 2) - idx.MustFragment("d", "general", 1).MustSetBit(11, SliceWidth+2) + idx.MustCreateFragmentIfNotExists("d", "general", 0).MustSetBit(11, 2) + idx.MustCreateFragmentIfNotExists("d", "general", 1).MustSetBit(11, SliceWidth+2) e := NewExecutor(idx.Index, NewCluster(1)) if res, err := e.Execute("d", MustParse(`union(get(id=10), get(id=11))`), nil); err != nil { @@ -99,9 +99,9 @@ func TestExecutor_Execute_Union(t *testing.T) { func TestExecutor_Execute_Count(t *testing.T) { idx := MustOpenIndex() defer idx.Close() - idx.MustFragment("d", "f", 0).MustSetBit(10, 3) - idx.MustFragment("d", "f", 1).MustSetBit(10, SliceWidth+1) - idx.MustFragment("d", "f", 1).MustSetBit(10, SliceWidth+2) + idx.MustCreateFragmentIfNotExists("d", "f", 0).MustSetBit(10, 3) + idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBit(10, SliceWidth+1) + idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBit(10, SliceWidth+2) e := NewExecutor(idx.Index, NewCluster(1)) if n, err := e.Execute("d", MustParse(`count(get(id=10, frame=f))`), nil); err != nil { @@ -121,7 +121,7 @@ func TestExecutor_Execute_Set(t *testing.T) { t.Fatal(err) } - f := idx.MustFragment("d", "f", 0) + f := idx.MustCreateFragmentIfNotExists("d", "f", 0) if n := f.Bitmap(10).Count(); n != 1 { t.Fatalf("unexpected bitmap count: %d", n) } @@ -142,7 +142,7 @@ func TestExecutor_Execute_Remote_Bitmap(t *testing.T) { t.Fatalf("unexpected db: %s", db) } else if query.String() != `get(id=10, frame=f)` { t.Fatalf("unexpected query: %s", query.String()) - } else if !reflect.DeepEqual(slices, []uint64{0, 2}) { + } else if !reflect.DeepEqual(slices, []uint64{0, 2, 4}) { t.Fatalf("unexpected slices: %+v", slices) } @@ -159,7 +159,7 @@ func TestExecutor_Execute_Remote_Bitmap(t *testing.T) { // The local node owns slice 1. idx := MustOpenIndex() defer idx.Close() - idx.MustFragment("d", "f", 1).MustSetBit(10, (1*SliceWidth)+1) + idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBit(10, (1*SliceWidth)+1) e := NewExecutor(idx.Index, c) if res, err := e.Execute("d", MustParse(`get(id=10, frame=f)`), nil); err != nil { @@ -190,8 +190,8 @@ func TestExecutor_Execute_Remote_Count(t *testing.T) { // Create local executor data. The local node owns slice 1. idx := MustOpenIndex() defer idx.Close() - idx.MustFragment("d", "f", 1).MustSetBit(10, (1*SliceWidth)+1) - idx.MustFragment("d", "f", 1).MustSetBit(10, (1*SliceWidth)+2) + idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBit(10, (1*SliceWidth)+1) + idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBit(10, (1*SliceWidth)+2) e := NewExecutor(idx.Index, c) if n, err := e.Execute("d", MustParse(`count(get(id=10, frame=f))`), nil); err != nil { @@ -233,7 +233,7 @@ func TestExecutor_Execute_Remote_Set(t *testing.T) { } // Verify that one bit is set on both node's index. - if n := idx.MustFragment("d", "f", 0).Bitmap(10).Count(); n != 1 { + if n := idx.MustCreateFragmentIfNotExists("d", "f", 0).Bitmap(10).Count(); n != 1 { t.Fatalf("unexpected local count: %d", n) } if !remoteCalled { diff --git a/fragment.go b/fragment.go index f16e9cd43..1157a6142 100644 --- a/fragment.go +++ b/fragment.go @@ -17,6 +17,9 @@ import ( // SliceWidth is the number of profile IDs in a slice. const SliceWidth = 65536 +// SnapshotExt is the file extension used for an in-process snapshot. +const SnapshotExt = ".snapshotting" + // Fragment represents the intersection of a frame and slice in a database. type Fragment struct { mu sync.Mutex @@ -64,61 +67,7 @@ func (f *Fragment) Open() error { defer f.mu.Unlock() // Initialize storage in a function so we can close if anything goes wrong. - if err := func() error { - // Create a roaring bitmap to serve as storage for the slice. - f.storage = roaring.NewBitmap() - - // Open the data file to be mmap'd and used as an ops log. - file, err := os.OpenFile(f.path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666) - if err != nil { - return fmt.Errorf("open file: %s", err) - } - f.file = file - - // Lock the underlying file. - if err := syscall.Flock(int(f.file.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err != nil { - return fmt.Errorf("flock: %s", err) - } - - // If the file is empty then initialize it with an empty bitmap. - fi, err := f.file.Stat() - if err != nil { - return err - } else if fi.Size() == 0 { - if _, err := f.storage.WriteTo(f.file); err != nil { - return fmt.Errorf("init storage file: %s", err) - } - - fi, err = f.file.Stat() - if err != nil { - return err - } - } - - // Mmap the underlying file so it can be zero copied. - storageData, err := syscall.Mmap(int(f.file.Fd()), 0, int(fi.Size()), syscall.PROT_READ, syscall.MAP_SHARED) - if err != nil { - return fmt.Errorf("mmap: %s", err) - } - f.storageData = storageData - - // Advise the kernel that the mmap is accessed randomly. - if err := madvise(f.storageData, syscall.MADV_RANDOM); err != nil { - return fmt.Errorf("madvise: %s", err) - } - - // Attach the mmap file to the bitmap. - data := (*[0x7FFFFFFF]byte)(unsafe.Pointer(&f.storageData[0]))[:fi.Size()] - if err := f.storage.UnmarshalBinary(data); err != nil { - return fmt.Errorf("unmarshal storage: file=%s, err=%s", f.file.Name(), err) - } - - // Attach the file to the bitmap to act as a write-ahead log. - f.storage.OpWriter = f.file - - return nil - - }(); err != nil { + if err := f.openStorage(); err != nil { f.close() return err } @@ -126,6 +75,63 @@ func (f *Fragment) Open() error { return nil } +// openStorage opens the storage bitmap. +func (f *Fragment) openStorage() error { + // Create a roaring bitmap to serve as storage for the slice. + f.storage = roaring.NewBitmap() + + // Open the data file to be mmap'd and used as an ops log. + file, err := os.OpenFile(f.path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666) + if err != nil { + return fmt.Errorf("open file: %s", err) + } + f.file = file + + // Lock the underlying file. + if err := syscall.Flock(int(f.file.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err != nil { + return fmt.Errorf("flock: %s", err) + } + + // If the file is empty then initialize it with an empty bitmap. + fi, err := f.file.Stat() + if err != nil { + return err + } else if fi.Size() == 0 { + if _, err := f.storage.WriteTo(f.file); err != nil { + return fmt.Errorf("init storage file: %s", err) + } + + fi, err = f.file.Stat() + if err != nil { + return err + } + } + + // Mmap the underlying file so it can be zero copied. + storageData, err := syscall.Mmap(int(f.file.Fd()), 0, int(fi.Size()), syscall.PROT_READ, syscall.MAP_SHARED) + if err != nil { + return fmt.Errorf("mmap: %s", err) + } + f.storageData = storageData + + // Advise the kernel that the mmap is accessed randomly. + if err := madvise(f.storageData, syscall.MADV_RANDOM); err != nil { + return fmt.Errorf("madvise: %s", err) + } + + // Attach the mmap file to the bitmap. + data := (*[0x7FFFFFFF]byte)(unsafe.Pointer(&f.storageData[0]))[:fi.Size()] + if err := f.storage.UnmarshalBinary(data); err != nil { + return fmt.Errorf("unmarshal storage: file=%s, err=%s", f.file.Name(), err) + } + + // Attach the file to the bitmap to act as a write-ahead log. + f.storage.OpWriter = f.file + + return nil + +} + // Close flushes the underlying storage, closes the file and unlocks it. func (f *Fragment) Close() error { f.mu.Lock() @@ -134,6 +140,13 @@ func (f *Fragment) Close() error { } func (f *Fragment) close() error { + if err := f.closeStorage(); err != nil { + return err + } + return nil +} + +func (f *Fragment) closeStorage() error { // Clear the storage bitmap so it doesn't access the closed mmap. f.storage = roaring.NewBitmap() @@ -234,7 +247,10 @@ func (f *Fragment) TopNAll(n int, categories []uint64) []Pair { func (f *Fragment) SetBit(bitmapID, profileID uint64) error { f.mu.Lock() defer f.mu.Unlock() + return f.setBit(bitmapID, profileID) +} +func (f *Fragment) setBit(bitmapID, profileID uint64) error { // Determine the position of the bit in the storage. pos, err := f.pos(bitmapID, profileID) if err != nil { @@ -420,6 +436,89 @@ func (f *Fragment) Range(bitmapID uint64, start, end time.Time) *Bitmap { return result } +// Import bulk imports a set of bits and then snapshots the storage. +// This does not affect the fragment's cache. +func (f *Fragment) Import(bitmapIDs, profileIDs []uint64) error { + f.mu.Lock() + defer f.mu.Unlock() + + // Verify that there are an equal number of bitmap ids and profile ids. + if len(bitmapIDs) != len(profileIDs) { + return fmt.Errorf("mismatch of bitmap and profile len: %d != %d", len(bitmapIDs), len(profileIDs)) + } + + // Disconnect op writer so we don't append updates. + f.storage.OpWriter = nil + + // Process every bit. + // If an error occurs then reopen the storage. + if err := func() error { + for i := range bitmapIDs { + // Determine the position of the bit in the storage. + pos, err := f.pos(bitmapIDs[i], profileIDs[i]) + if err != nil { + return err + } + + // Write to storage. + if err := f.storage.Add(pos); err != nil { + return err + } + } + return nil + }(); err != nil { + _ = f.closeStorage() + _ = f.openStorage() + return err + } + + // Write the storage to disk and reload. + if err := f.snapshot(); err != nil { + return err + } + + return nil +} + +// Snapshot writes the storage bitmap to disk and reopens it. +func (f *Fragment) Snapshot() error { + f.mu.Lock() + defer f.mu.Unlock() + return f.snapshot() +} + +func (f *Fragment) snapshot() error { + // Create a temporary file to snapshot to. + snapshotPath := f.path + SnapshotExt + file, err := os.Create(snapshotPath) + if err != nil { + return fmt.Errorf("create snapshot file: %s", err) + } + defer file.Close() + + // Write storage to snapshot. + if _, err := f.storage.WriteTo(file); err != nil { + return fmt.Errorf("snapshot write to: %s", err) + } + + // Close current storage. + if err := f.closeStorage(); err != nil { + return fmt.Errorf("close storage: %s", err) + } + + // Move snapshot to data file location. + if err := os.Rename(snapshotPath, f.path); err != nil { + return fmt.Errorf("rename snapshot: %s", err) + } + + // Reopen storage. + if err := f.openStorage(); err != nil { + return fmt.Errorf("open storage: %s", err) + } + + return nil +} + func madvise(b []byte, advice int) (err error) { _, _, e1 := syscall.Syscall(syscall.SYS_MADVISE, uintptr(unsafe.Pointer(&b[0])), uintptr(len(b)), uintptr(advice)) if e1 != 0 { diff --git a/fragment_test.go b/fragment_test.go index 306914f82..ad8057be9 100644 --- a/fragment_test.go +++ b/fragment_test.go @@ -69,6 +69,35 @@ func TestFragment_ClearBit(t *testing.T) { } } +// Ensure a fragment can snapshot correctly. +func TestFragment_Snapshot(t *testing.T) { + f := MustOpenFragment("d", "f", 0) + defer f.Close() + + // Set and then clear bits on the fragment. + if err := f.SetBit(1000, 1); err != nil { + t.Fatal(err) + } else if err := f.SetBit(1000, 2); err != nil { + t.Fatal(err) + } else if err := f.ClearBit(1000, 1); err != nil { + t.Fatal(err) + } + + // Snapshot bitmap and verify data. + if err := f.Snapshot(); err != nil { + t.Fatal(err) + } else if n := f.Bitmap(1000).Count(); n != 1 { + t.Fatalf("unexpected count: %d", n) + } + + // Close and reopen the fragment & verify the data. + if err := f.Reopen(); err != nil { + t.Fatal(err) + } else if n := f.Bitmap(1000).Count(); n != 1 { + t.Fatalf("unexpected count (reopen): %d", n) + } +} + // Fragment is a test wrapper for pilosa.Fragment. type Fragment struct { *pilosa.Fragment diff --git a/handler.go b/handler.go index fa37372b7..60b99ba53 100644 --- a/handler.go +++ b/handler.go @@ -20,6 +20,12 @@ import ( // Handler represents an HTTP handler. type Handler struct { + Index *Index + + // Local hostname & cluster configuration. + Host string + Cluster *Cluster + // The execution engine for running queries. Executor interface { Execute(db string, query *pql.Query, slices []uint64) (interface{}, error) @@ -44,7 +50,19 @@ func NewHandler() *Handler { func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/query": - h.handleQuery(w, r) + switch r.Method { + case "POST": + h.handlePostQuery(w, r) + default: + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + } + case "/import": + switch r.Method { + case "POST": + h.handlePostImport(w, r) + default: + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + } case "/version": h.handleVersion(w, r) @@ -55,16 +73,6 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } } -// handleQuery handles /query requests. -func (h *Handler) handleQuery(w http.ResponseWriter, r *http.Request) { - switch r.Method { - case "POST": - h.handlePostQuery(w, r) - default: - http.Error(w, "method not allowed", http.StatusMethodNotAllowed) - } -} - // handlePostQuery handles /query requests. func (h *Handler) handlePostQuery(w http.ResponseWriter, r *http.Request) { // Parse incoming request. @@ -212,6 +220,66 @@ func (h *Handler) writeJSONQueryResponse(w http.ResponseWriter, res interface{}, return json.NewEncoder(w).Encode(o) } +// handlePostImport handles /import requests. +func (h *Handler) handlePostImport(w http.ResponseWriter, r *http.Request) { + // Verify that request is only communicating over protobufs. + if r.Header.Get("Content-Type") != "application/x-protobuf" { + http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType) + return + } else if r.Header.Get("Accept") != "application/x-protobuf" { + http.Error(w, "Not acceptable", http.StatusNotAcceptable) + return + } + + // Read entire body. + body, err := ioutil.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + // Marshal into request object. + var req internal.ImportRequest + if err := proto.Unmarshal(body, &req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + db, frame, slice := req.GetDB(), req.GetFrame(), req.GetSlice() + + // Validate that this handler owns the slice. + if !h.Cluster.OwnsSlice(h.Host, slice) { + http.Error(w, "host does not own slice", http.StatusPreconditionFailed) + return + } + + // Find the correct fragment. + f, err := h.Index.CreateFragmentIfNotExists(db, frame, slice) + if err != nil { + h.logger().Printf("fragment error: db=%s, frame=%s, slice=%s, err=%s", db, frame, slice, err) + http.Error(w, "fragment error", http.StatusInternalServerError) + return + } + + // Import into fragment. + err = f.Import(req.GetBitmapIDs(), req.GetProfileIDs()) + if err != nil { + h.logger().Printf("import error: db=%s, frame=%s, slice=%s, bits=%d", db, frame, slice, len(req.GetProfileIDs())) + } + + // Marshal response object. + buf, e := proto.Marshal(&internal.ImportResponse{Err: proto.String(errorString(err))}) + if e != nil { + http.Error(w, fmt.Sprintf("marshal import response: %s", err), http.StatusInternalServerError) + return + } + + // Write response. + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + } + w.Write(buf) +} + // handleGetVersion handles /version requests. func (h *Handler) handleVersion(w http.ResponseWriter, r *http.Request) { if err := json.NewEncoder(w).Encode(struct { @@ -262,3 +330,11 @@ func parseUint64Slice(s string) ([]uint64, error) { } return a, nil } + +// errorString returns the string representation of err. +func errorString(err error) string { + if err == nil { + return "" + } + return err.Error() +} diff --git a/handler_test.go b/handler_test.go index f2fb0c2d8..33db15b7e 100644 --- a/handler_test.go +++ b/handler_test.go @@ -145,7 +145,7 @@ func TestHandler_Query_Bitmap_JSON(t *testing.T) { h.ServeHTTP(w, MustNewHTTPRequest("POST", "/query", strings.NewReader("get(100)"))) if w.Code != http.StatusOK { t.Fatalf("unexpected status code: %d", w.Code) - } else if body := w.Body.String(); body != `{"result":{"chunks":[{"Key":0,"Value":[10,4,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]},{"Key":32,"Value":[2,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]}]}}`+"\n" { + } else if body := w.Body.String(); body != `{"result":[1,3,66,65537]}`+"\n" { t.Fatalf("unexpected body: %q", body) } } diff --git a/index.go b/index.go index c2025fc84..5bf25de59 100644 --- a/index.go +++ b/index.go @@ -30,6 +30,83 @@ func (i *Index) Open() error { if err := os.MkdirAll(i.path, 0777); err != nil { return err } + + // Open all databases. + if err := i.openDatabases(); err != nil { + return err + } + + return nil +} + +// openDatabases recursively opens all directories within the data directory. +func (i *Index) openDatabases() error { + f, err := os.Open(i.path) + if err != nil { + return err + } + defer f.Close() + + fis, err := f.Readdir(0) + if err != nil { + return err + } + + for _, fi := range fis { + if !fi.IsDir() { + continue + } else if err := i.openDatabase(filepath.Base(fi.Name())); err != nil { + return err + } + } + return nil +} + +// openDatabase recursively opens all frames within a database directory. +func (i *Index) openDatabase(db string) error { + f, err := os.Open(filepath.Join(i.path, db)) + if err != nil { + return err + } + defer f.Close() + + fis, err := f.Readdir(0) + if err != nil { + return err + } + + for _, fi := range fis { + if !fi.IsDir() { + continue + } else if err := i.openFrame(db, filepath.Base(fi.Name())); err != nil { + return err + } + } + return nil +} + +// openFrame recursively opens all fragments within a frame directory. +func (i *Index) openFrame(db, frame string) error { + f, err := os.Open(filepath.Join(i.path, db, frame)) + if err != nil { + return err + } + defer f.Close() + + fis, err := f.Readdir(0) + if err != nil { + return err + } + + for _, fi := range fis { + slice, err := strconv.ParseUint(filepath.Base(fi.Name()), 10, 64) + if err != nil || fi.IsDir() { + continue + } + if _, err := i.CreateFragmentIfNotExists(db, frame, slice); err != nil { + return fmt.Errorf("open fragment: db=%s, frame=%d, slice=%d, err=%s", db, frame, slice, err) + } + } return nil } @@ -59,8 +136,15 @@ func (i *Index) FragmentPath(db, frame string, slice uint64) string { } // Fragment returns the fragment for a database, frame & slice. +func (i *Index) Fragment(db, frame string, slice uint64) *Fragment { + i.mu.Lock() + defer i.mu.Unlock() + return i.fragments[fragmentKey{db, frame, slice}] +} + +// CreateFragmentIfNotExists returns the fragment for a database, frame & slice. // The fragment is created if it doesn't already exist. -func (i *Index) Fragment(db, frame string, slice uint64) (*Fragment, error) { +func (i *Index) CreateFragmentIfNotExists(db, frame string, slice uint64) (*Fragment, error) { i.mu.Lock() defer i.mu.Unlock() diff --git a/index_test.go b/index_test.go index 4f4d66597..440416d47 100644 --- a/index_test.go +++ b/index_test.go @@ -36,9 +36,9 @@ func (i *Index) Close() error { return i.Index.Close() } -// MustFragment returns a given fragment. Panic on error. -func (i *Index) MustFragment(db, frame string, slice uint64) *Fragment { - f, err := i.Index.Fragment(db, frame, slice) +// MustCreateFragmentIfNotExists returns a given fragment. Panic on error. +func (i *Index) MustCreateFragmentIfNotExists(db, frame string, slice uint64) *Fragment { + f, err := i.Index.CreateFragmentIfNotExists(db, frame, slice) if err != nil { panic(err) } diff --git a/internal/internal.pb.go b/internal/internal.pb.go index bdb18ac33..0a557144b 100644 --- a/internal/internal.pb.go +++ b/internal/internal.pb.go @@ -12,8 +12,11 @@ It has these top-level messages: Bitmap Chunk Pair + Bit QueryRequest QueryResponse + ImportRequest + ImportResponse */ package internal @@ -88,6 +91,30 @@ func (m *Pair) GetCount() uint64 { return 0 } +type Bit struct { + BitmapID *uint64 `protobuf:"varint,1,req" json:"BitmapID,omitempty"` + ProfileID *uint64 `protobuf:"varint,2,req" json:"ProfileID,omitempty"` + XXX_unrecognized []byte `json:"-"` +} + +func (m *Bit) Reset() { *m = Bit{} } +func (m *Bit) String() string { return proto.CompactTextString(m) } +func (*Bit) ProtoMessage() {} + +func (m *Bit) GetBitmapID() uint64 { + if m != nil && m.BitmapID != nil { + return *m.BitmapID + } + return 0 +} + +func (m *Bit) GetProfileID() uint64 { + if m != nil && m.ProfileID != nil { + return *m.ProfileID + } + return 0 +} + type QueryRequest struct { DB *string `protobuf:"bytes,1,req" json:"DB,omitempty"` Query *string `protobuf:"bytes,2,req" json:"Query,omitempty"` @@ -160,5 +187,69 @@ func (m *QueryResponse) GetPairs() []*Pair { return nil } +type ImportRequest struct { + DB *string `protobuf:"bytes,1,req" json:"DB,omitempty"` + Frame *string `protobuf:"bytes,2,req" json:"Frame,omitempty"` + Slice *uint64 `protobuf:"varint,3,req" json:"Slice,omitempty"` + BitmapIDs []uint64 `protobuf:"varint,4,rep" json:"BitmapIDs,omitempty"` + ProfileIDs []uint64 `protobuf:"varint,5,rep" json:"ProfileIDs,omitempty"` + XXX_unrecognized []byte `json:"-"` +} + +func (m *ImportRequest) Reset() { *m = ImportRequest{} } +func (m *ImportRequest) String() string { return proto.CompactTextString(m) } +func (*ImportRequest) ProtoMessage() {} + +func (m *ImportRequest) GetDB() string { + if m != nil && m.DB != nil { + return *m.DB + } + return "" +} + +func (m *ImportRequest) GetFrame() string { + if m != nil && m.Frame != nil { + return *m.Frame + } + return "" +} + +func (m *ImportRequest) GetSlice() uint64 { + if m != nil && m.Slice != nil { + return *m.Slice + } + return 0 +} + +func (m *ImportRequest) GetBitmapIDs() []uint64 { + if m != nil { + return m.BitmapIDs + } + return nil +} + +func (m *ImportRequest) GetProfileIDs() []uint64 { + if m != nil { + return m.ProfileIDs + } + return nil +} + +type ImportResponse struct { + Err *string `protobuf:"bytes,1,opt" json:"Err,omitempty"` + XXX_unrecognized []byte `json:"-"` +} + +func (m *ImportResponse) Reset() { *m = ImportResponse{} } +func (m *ImportResponse) String() string { return proto.CompactTextString(m) } +func (*ImportResponse) ProtoMessage() {} + +func (m *ImportResponse) GetErr() string { + if m != nil && m.Err != nil { + return *m.Err + } + return "" +} + func init() { } diff --git a/internal/internal.proto b/internal/internal.proto index 9630bc798..ac1b97bca 100644 --- a/internal/internal.proto +++ b/internal/internal.proto @@ -14,6 +14,11 @@ message Pair { required uint64 Count = 2; } +message Bit { + required uint64 BitmapID = 1; + required uint64 ProfileID = 2; +} + message QueryRequest { required string DB = 1; required string Query = 2; @@ -26,3 +31,15 @@ message QueryResponse { optional uint64 N = 3; repeated Pair Pairs = 4; } + +message ImportRequest { + required string DB = 1; + required string Frame = 2; + required uint64 Slice = 3; + repeated uint64 BitmapIDs = 4; + repeated uint64 ProfileIDs = 5; +} + +message ImportResponse { + optional string Err = 1; +} diff --git a/pilosa.go b/pilosa.go index 19c291ffd..89e672f81 100644 --- a/pilosa.go +++ b/pilosa.go @@ -1,6 +1,21 @@ package pilosa +import ( + "errors" +) + //go:generate protoc --gogo_out=. internal/internal.proto +var ( + // ErrHostRequired is returned when excuting a remote operation without a host. + ErrHostRequired = errors.New("host required") + + // ErrDatabaseRequired is returned when no database is specified. + ErrDatabaseRequired = errors.New("database required") + + // ErrFrameRequired is returned when no frame is specified. + ErrFrameRequired = errors.New("frame required") +) + // Version represents the current running version of Pilosa. var Version string