Merge pull request #32 from benbjohnson/cluster-import

Cluster-wide Import

client.go

@@ -2,6 +2,7 @@ package pilosa
 
 import (
 	"bytes"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -37,6 +38,31 @@ func NewClient(host string) (*Client, error) {
 // Host returns the host the client was initialized with.
 func (c *Client) Host() string { return c.host }
 
+// SliceNodes returns a list of nodes that own a slice.
+func (c *Client) SliceNodes(slice uint64) ([]*Node, error) {
+	// Execute request against the host.
+	u := url.URL{
+		Scheme:   "http",
+		Host:     c.host,
+		Path:     "/slices/nodes",
+		RawQuery: (url.Values{"slice": {strconv.FormatUint(slice, 10)}}).Encode(),
+	}
+	resp, err := c.HTTPClient.Get(u.String())
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	var a []*Node
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("http: status=%d", resp.StatusCode)
+	} else if err := json.NewDecoder(resp.Body).Decode(&a); err != nil {
+		return nil, fmt.Errorf("json decode: %s", err)
+	}
+
+	return a, nil
+}
+
 // 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 == "" {
@@ -45,6 +71,12 @@ func (c *Client) Import(db, frame string, slice uint64, bits []Bit) error {
 		return ErrFrameRequired
 	}
 
+	// Retrieve a list of nodes that own the slice.
+	nodes, err := c.SliceNodes(slice)
+	if err != nil {
+		return fmt.Errorf("slice nodes: %s", err)
+	}
+
 	// Separate bitmap and profile IDs to reduce allocations.
 	bitmapIDs := Bits(bits).BitmapIDs()
 	profileIDs := Bits(bits).ProfileIDs()
@@ -61,8 +93,20 @@ func (c *Client) Import(db, frame string, slice uint64, bits []Bit) error {
 		return fmt.Errorf("marshal import request: %s", err)
 	}
 
+	// Import to each node.
+	for _, node := range nodes {
+		if err := c.importNode(node, buf); err != nil {
+			return fmt.Errorf("import node: host=%s, err=%s", node.Host, err)
+		}
+	}
+
+	return nil
+}
+
+// importNode sends a pre-marshaled import request to a node.
+func (c *Client) importNode(node *Node, buf []byte) error {
 	// Create URL & HTTP request.
-	u := url.URL{Scheme: "http", Host: c.host, Path: "/import"}
+	u := url.URL{Scheme: "http", Host: node.Host, Path: "/import"}
 	req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf))
 	if err != nil {
 		return err

cluster.go

@@ -15,7 +15,7 @@ const (
 
 // Node represents a node in the cluster.
 type Node struct {
-	Host string
+	Host string `json:"host"`
 }
 
 // Nodes represents a list of nodes.

cmd/pilosa/main.go

@@ -13,6 +13,7 @@ import (
 	"path/filepath"
 	"runtime/pprof"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/BurntSushi/toml"
@@ -31,6 +32,9 @@ func init() {
 }
 
 const (
+	// DefaultDataDir is the default data directory.
+	DefaultDataDir = "~/.pilosa"
+
 	// DefaultHost is the default hostname and port to use.
 	DefaultHost = "localhost:15000"
 )
@@ -196,15 +200,21 @@ func (m *Main) ParseFlags(args []string) error {
 		}
 	}
 
-	// If no data directory is specified then use ~/.pilosa
+	// Use default data directory if one is not specified.
 	if m.Config.DataDir == "" {
+		m.Config.DataDir = DefaultDataDir
+	}
+
+	// Expand home directory.
+	prefix := "~" + string(filepath.Separator)
+	if strings.HasPrefix(m.Config.DataDir, prefix) {
 		u, err := user.Current()
 		if err != nil {
 			return err
 		} else if u.HomeDir == "" {
 			return errors.New("data directory not specified and no home dir available")
 		}
-		m.Config.DataDir = filepath.Join(u.HomeDir, ".pilosa")
+		m.Config.DataDir = filepath.Join(u.HomeDir, strings.TrimPrefix(m.Config.DataDir, prefix))
 	}
 
 	return nil
@@ -217,7 +227,7 @@ type Config struct {
 
 	Cluster struct {
 		ReplicaN int           `toml:"replicas"`
-		Nodes    []*ConfigNode `toml:"nodes"`
+		Nodes    []*ConfigNode `toml:"node"`
 	} `toml:"cluster"`
 
 	Plugins struct {

cmd/pilosactl/main.go

@@ -74,6 +74,8 @@ func (m *Main) Run() error {
 	switch m.Command {
 	case "", "help", "-h":
 		return ErrUsage
+	case "config":
+		cmd = NewConfigCommand(m.Stdin, m.Stdout, m.Stderr)
 	case "import":
 		cmd = NewImportCommand(m.Stdin, m.Stdout, m.Stderr)
 	default:
@@ -113,6 +115,60 @@ type Command interface {
 	Run() error
 }
 
+// ConfigCommand represents a command for printing a default config.
+type ConfigCommand struct {
+	// Standard input/output
+	Stdin  io.Reader
+	Stdout io.Writer
+	Stderr io.Writer
+}
+
+// NewConfigCommand returns a new instance of ConfigCommand.
+func NewConfigCommand(stdin io.Reader, stdout, stderr io.Writer) *ConfigCommand {
+	return &ConfigCommand{
+		Stdin:  stdin,
+		Stdout: stdout,
+		Stderr: stderr,
+	}
+}
+
+// ParseFlags parses command line flags from args.
+func (cmd *ConfigCommand) ParseFlags(args []string) error {
+	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
+	fs.SetOutput(cmd.Stderr)
+	if err := fs.Parse(args); err != nil {
+		return err
+	}
+	return nil
+}
+
+// Usage returns the usage message to be printed.
+func (cmd *ConfigCommand) Usage() string {
+	return strings.TrimSpace(`
+usage: pilosactl config
+
+Prints the default configuration file to standard out.
+`)
+}
+
+// Run executes the main program execution.
+func (cmd *ConfigCommand) Run() error {
+	fmt.Fprintln(cmd.Stdout, strings.TrimSpace(`
+data-dir = "~/.pilosa"
+host = "localhost:15000"
+
+[cluster]
+replicas = 1
+
+[[cluster.node]]
+host = "localhost:15000"
+
+[plugins]
+path = ""
+`)+"\n")
+	return nil
+}
+
 // ImportCommand represents a command for bulk importing data.
 type ImportCommand struct {
 	// Destination host and port.

handler.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/umbel/pilosa/internal"
@@ -48,6 +49,8 @@ func NewHandler() *Handler {
 
 // ServeHTTP handles an HTTP request.
 func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	t := time.Now()
+
 	switch r.URL.Path {
 	case "/query":
 		switch r.Method {
@@ -63,6 +66,13 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		default:
 			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
 		}
+	case "/slices/nodes":
+		switch r.Method {
+		case "GET":
+			h.handleGetSlicesNodes(w, r)
+		default:
+			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+		}
 	case "/version":
 		h.handleVersion(w, r)
 
@@ -71,6 +81,8 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	default:
 		http.NotFound(w, r)
 	}
+
+	h.logger().Printf("%s %s %.03fs", r.Method, r.URL.String(), time.Since(t).Seconds())
 }
 
 // handlePostQuery handles /query requests.
@@ -280,6 +292,25 @@ func (h *Handler) handlePostImport(w http.ResponseWriter, r *http.Request) {
 	w.Write(buf)
 }
 
+// handleGetSlicesNodes handles /slices/nodes requests.
+func (h *Handler) handleGetSlicesNodes(w http.ResponseWriter, r *http.Request) {
+	q := r.URL.Query()
+
+	// Read slice parameter.
+	slice, err := strconv.ParseUint(q.Get("slice"), 10, 64)
+	if err != nil {
+		http.Error(w, "slice required", http.StatusBadRequest)
+	}
+
+	// Retrieve slice owner nodes.
+	nodes := h.Cluster.SliceNodes(slice)
+
+	// Write to response.
+	if err := json.NewEncoder(w).Encode(nodes); err != nil {
+		h.logger().Printf("json write error: %s", err)
+	}
+}
+
 // handleGetVersion handles /version requests.
 func (h *Handler) handleVersion(w http.ResponseWriter, r *http.Request) {
 	if err := json.NewEncoder(w).Encode(struct {

handler_test.go

@@ -293,6 +293,22 @@ func TestHandler_Version(t *testing.T) {
 	}
 }
 
+// Ensure the handler can return a list of nodes for a slice.
+func TestHandler_Slices_Nodes(t *testing.T) {
+	h := NewHandler()
+	h.Cluster = NewCluster(3)
+	h.Cluster.ReplicaN = 2
+
+	w := httptest.NewRecorder()
+	r := MustNewHTTPRequest("GET", "/slices/nodes?slice=0", nil)
+	h.ServeHTTP(w, r)
+	if w.Code != http.StatusOK {
+		t.Fatalf("unexpected status code: %d", w.Code)
+	} else if w.Body.String() != `[{"host":"host2"},{"host":"host0"}]`+"\n" {
+		t.Fatalf("unexpected body: %q", w.Body.String())
+	}
+}
+
 // Ensure the handler can return expvars without panicking.
 func TestHandler_Expvars(t *testing.T) {
 	h := NewHandler()