add frame restore for cluster transfer

This commit adds an endpoint to restore an entire frame from
another cluster. Multiple hosts can use this endpoint to copy
and rebalance a cluster to a new cluster.

client.go

@@ -203,30 +203,21 @@ func (c *Client) BackupTo(w io.Writer, db, frame string) error {
 
 // 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)
+	r, err := c.BackupSlice(db, frame, slice)
+	if err != nil {
+		return fmt.Errorf("backup slice: slice=%d, err=%s", slice, err)
+	} else if r == nil {
+		return nil
+	}
+	defer r.Close()
+
+	// Read entire buffer to determine file size.
+	data, err := ioutil.ReadAll(r)
+	if err != nil {
+		return err
+	} else if err := r.Close(); err != nil {
+		return err
 	}
 
 	// Write slice file header.
@@ -247,7 +238,32 @@ func (c *Client) backupSliceTo(tw *tar.Writer, db, frame string, slice uint64) e
 	return nil
 }
 
-func (c *Client) backupSliceNode(db, frame string, slice uint64, node *Node) ([]byte, error) {
+// BackupSlice retrieves a streaming backup from a single slice.
+// This function tries slice owners until one succeeds.
+func (c *Client) BackupSlice(db, frame string, slice uint64) (io.ReadCloser, error) {
+	// Retrieve a list of nodes that own the slice.
+	nodes, err := c.SliceNodes(slice)
+	if err != nil {
+		return nil, fmt.Errorf("slice nodes: %s", err)
+	}
+
+	// Try to backup slice from each one until successful.
+	for _, i := range rand.Perm(len(nodes)) {
+		r, err := c.backupSliceNode(db, frame, slice, nodes[i])
+		if err == nil {
+			return r, nil // successfully attached
+		} else if err == ErrFragmentNotFound {
+			return nil, nil // slice doesn't exist
+		} else if err != nil {
+			log.Println(err)
+			continue
+		}
+	}
+
+	return nil, fmt.Errorf("unable to connect to any owner")
+}
+
+func (c *Client) backupSliceNode(db, frame string, slice uint64, node *Node) (io.ReadCloser, error) {
 	u := url.URL{
 		Scheme: "http",
 		Host:   node.Host,
@@ -262,16 +278,17 @@ func (c *Client) backupSliceNode(db, frame string, slice uint64, node *Node) ([]
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
 
 	// Return error if status is not OK.
 	if resp.StatusCode == http.StatusNotFound {
+		resp.Body.Close()
 		return nil, ErrFragmentNotFound
 	} else if resp.StatusCode != http.StatusOK {
+		resp.Body.Close()
 		return nil, fmt.Errorf("unexpected backup status code: host=%s, code=%d", node.Host, resp.StatusCode)
 	}
 
-	return ioutil.ReadAll(resp.Body)
+	return resp.Body, nil
 }
 
 // RestoreFrom restores a frame from a backup file to an entire cluster.
@@ -348,6 +365,32 @@ func (c *Client) restoreSliceFrom(buf []byte, db, frame string, slice uint64) er
 	return nil
 }
 
+// RestoreFrame restores an entire frame from a host in another cluster.
+func (c *Client) RestoreFrame(host, db, frame string) error {
+	u := url.URL{
+		Scheme: "http",
+		Host:   c.Host(),
+		Path:   "/frame/restore",
+		RawQuery: url.Values{
+			"host":  {host},
+			"db":    {db},
+			"frame": {frame},
+		}.Encode(),
+	}
+	resp, err := c.HTTPClient.Post(u.String(), "application/octet-stream", nil)
+	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", host, resp.StatusCode)
+	}
+
+	return nil
+}
+
 // Bit represents the location of a single bit.
 type Bit struct {
 	BitmapID  uint64

cluster.go

@@ -7,7 +7,7 @@ import (
 
 const (
 	// DefaultPartitionN is the default number of partitions in a cluster.
-	DefaultPartitionN = 64
+	DefaultPartitionN = 16
 
 	// DefaultReplicaN is the default number of replicas per partition.
 	DefaultReplicaN = 1
@@ -113,6 +113,9 @@ type Hasher interface {
 	Hash(key uint64, n int) int
 }
 
+// NewHasher returns a new instance of the default hasher.
+func NewHasher() Hasher { return &jmphasher{} }
+
 // jmphasher represents an implementation of jmphash. Implements Hasher.
 type jmphasher struct{}
 

cluster_test.go

@@ -56,6 +56,26 @@ func TestCluster_Partition(t *testing.T) {
 	}
 }
 
+// Ensure the hasher can hash correctly.
+func TestHasher(t *testing.T) {
+	for _, tt := range []struct {
+		key    uint64
+		bucket []int
+	}{
+		// Generated from the reference C++ code
+		{0, []int{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}},
+		{1, []int{0, 0, 0, 0, 0, 0, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 17, 17}},
+		{0xdeadbeef, []int{0, 1, 2, 3, 3, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 16, 16, 16}},
+		{0x0ddc0ffeebadf00d, []int{0, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 15, 15, 15, 15}},
+	} {
+		for i, v := range tt.bucket {
+			if got := pilosa.NewHasher().Hash(tt.key, i+1); got != v {
+				t.Errorf("hash(%v,%v)=%v, want %v", tt.key, i+1, got, v)
+			}
+		}
+	}
+}
+
 // NewCluster returns a cluster with n nodes and uses a mod-based hasher.
 func NewCluster(n int) *pilosa.Cluster {
 	c := pilosa.NewCluster()

cmd/pilosa/main.go

@@ -87,6 +87,10 @@ type Main struct {
 	// Configuration options.
 	Config *Config
 
+	// Cluster configuration shared by components
+	Host    string
+	Cluster *pilosa.Cluster
+
 	// Profiling paths
 	CPUProfile string
 
@@ -148,13 +152,13 @@ func (m *Main) Run(args ...string) error {
 	m.ln = ln
 
 	// Determine hostname based on listening port.
-	hostname := net.JoinHostPort(host, strconv.Itoa(m.ln.Addr().(*net.TCPAddr).Port))
+	m.Host = net.JoinHostPort(host, strconv.Itoa(m.ln.Addr().(*net.TCPAddr).Port))
 
 	// Build cluster from config file. Create local host if none are specified.
-	cluster := m.Config.PilosaCluster()
-	if len(cluster.Nodes) == 0 {
-		cluster.Nodes = []*pilosa.Node{{
-			Host: hostname,
+	m.Cluster = m.Config.PilosaCluster()
+	if len(m.Cluster.Nodes) == 0 {
+		m.Cluster.Nodes = []*pilosa.Node{{
+			Host: m.Host,
 		}}
 	}
 
@@ -167,43 +171,43 @@ func (m *Main) Run(args ...string) error {
 
 	// Create executor for executing queries.
 	e := pilosa.NewExecutor(m.index)
-	e.Host = hostname
-	e.Cluster = cluster
+	e.Host = m.Host
+	e.Cluster = m.Cluster
 
 	// Initialize HTTP handler.
 	h := pilosa.NewHandler()
 	h.Index = m.index
-	h.Host = hostname
-	h.Cluster = cluster
+	h.Host = m.Host
+	h.Cluster = m.Cluster
 	h.Executor = e
 	h.LogOutput = m.Stderr
 
 	// Serve HTTP.
 	go func() { http.Serve(ln, h) }()
 
-	//sync up max slice if more than one node
-	if len(cluster.Nodes) > 1 { 
+	// Sync up max slice if more than one node
+	if len(m.Cluster.Nodes) > 1 {
 		m.ticker = time.NewTicker(time.Second * time.Duration(m.pollingSecs))
-		go func() { 
+		go func() {
 			for range m.ticker.C {
-				oldmax:= m.index.SliceN()
-				newmax:=oldmax
-				for _, node := range cluster.Nodes {
-					if hostname != node.Host {
-						newslice,_:=checkMaxSlice(node.Host)
-							if newslice>newmax{
-								newmax= newslice
-							}
+				oldmax := m.index.SliceN()
+				newmax := oldmax
+				for _, node := range m.Cluster.Nodes {
+					if m.Host != node.Host {
+						newslice, _ := checkMaxSlice(node.Host)
+						if newslice > newmax {
+							newmax = newslice
+						}
 					}
 				}
-				if newmax>oldmax{
+				if newmax > oldmax {
 					m.index.SetMax(newmax)
 				}
 			}
 		}()
 	}
 
-	fmt.Fprintf(m.Stderr, "Listening as http://%s\n", hostname)
+	fmt.Fprintf(m.Stderr, "Listening as http://%s\n", m.Host)
 
 	return nil
 }

cmd/pilosa/main_test.go

@@ -16,6 +16,7 @@ import (
 	"testing/quick"
 
 	"github.com/BurntSushi/toml"
+	"github.com/umbel/pilosa"
 	main "github.com/umbel/pilosa/cmd/pilosa"
 )
 
@@ -173,6 +174,61 @@ func TestMain_SetProfileAttrs(t *testing.T) {
 	}
 }
 
+// Ensure program can set bits on one cluster and then restore to a second cluster.
+func TestMain_FrameRestore(t *testing.T) {
+	m0 := MustRunMain()
+	defer m0.Close()
+
+	m1 := MustRunMain()
+	defer m1.Close()
+
+	// Update cluster config.
+	m0.Cluster.Nodes = []*pilosa.Node{
+		{Host: m0.Host},
+		{Host: m1.Host},
+	}
+	m1.Cluster.Nodes = m0.Cluster.Nodes
+
+	// Write data on first cluster.
+	if _, err := m0.Query("db=d", `
+		SetBit(id=1, frame="f", profileID=100)
+		SetBit(id=1, frame="f", profileID=1000)
+		SetBit(id=1, frame="f", profileID=100000)
+		SetBit(id=1, frame="f", profileID=200000)
+		SetBit(id=1, frame="f", profileID=400000)
+		SetBit(id=1, frame="f", profileID=600000)
+		SetBit(id=1, frame="f", profileID=800000)
+	`); err != nil {
+		t.Fatal(err)
+	}
+
+	// Query bitmap on first cluster.
+	if res, err := m0.Query("db=d", `Bitmap(id=1, frame="f")`); err != nil {
+		t.Fatal(err)
+	} else if res != `{"results":[{"attrs":{},"bits":[100,1000,100000,200000,400000,600000,800000]}]}`+"\n" {
+		t.Fatalf("unexpected result: %s", res)
+	}
+
+	// Start second cluster.
+	m2 := MustRunMain()
+	defer m2.Close()
+
+	// Import from first cluster.
+	client, err := pilosa.NewClient(m2.Host)
+	if err != nil {
+		t.Fatal(err)
+	} else if err := client.RestoreFrame(m0.Host, "d", "f"); err != nil {
+		t.Fatal(err)
+	}
+
+	// Query bitmap on second cluster.
+	if res, err := m2.Query("db=d", `Bitmap(id=1, frame="f")`); err != nil {
+		t.Fatal(err)
+	} else if res != `{"results":[{"attrs":{},"bits":[100,1000,100000,200000,400000,600000,800000]}]}`+"\n" {
+		t.Fatalf("unexpected result: %s", res)
+	}
+}
+
 // Ensure the host can be parsed.
 func TestConfig_Parse_Host(t *testing.T) {
 	if c, err := ParseConfig(`host = "local"`); err != nil {

executor.go

@@ -402,6 +402,7 @@ func (e *Executor) executeClearBit(db string, c *pql.ClearBit, opt *ExecOptions)
 func (e *Executor) executeSetBit(db string, c *pql.SetBit, opt *ExecOptions) (bool, error) {
 	slice := c.ProfileID / SliceWidth
 	ret := false
+
 	for _, node := range e.Cluster.SliceNodes(slice) {
 		// Update locally if host matches.
 		if node.Host == e.Host {

handler.go

@@ -89,6 +89,13 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		default:
 			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
 		}
+	case "/frame/restore":
+		switch r.Method {
+		case "POST":
+			h.handlePostFrameRestore(w, r)
+		default:
+			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+		}
 	case "/version":
 		h.handleVersion(w, r)
 
@@ -435,6 +442,75 @@ func (h *Handler) handlePostFragmentData(w http.ResponseWriter, r *http.Request)
 	}
 }
 
+// handlePostFrameRestore handles POST /frame/restore requests.
+func (h *Handler) handlePostFrameRestore(w http.ResponseWriter, r *http.Request) {
+	q := r.URL.Query()
+	host := q.Get("host")
+	db, frame := q.Get("db"), q.Get("frame")
+
+	// Validate query parameters.
+	if host == "" {
+		http.Error(w, "host required", http.StatusBadRequest)
+		return
+	} else if db == "" {
+		http.Error(w, "db required", http.StatusBadRequest)
+		return
+	} else if frame == "" {
+		http.Error(w, "frame required", http.StatusBadRequest)
+		return
+	}
+
+	// Create a client for the remote cluster.
+	client, err := NewClient(host)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	// Determine the maximum number of slices.
+	sliceN, err := client.SliceN()
+	if err != nil {
+		http.Error(w, "cannot determine remote slice count: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	// Loop over each slice and import it if this node owns it.
+	for slice := uint64(0); slice <= sliceN; slice++ {
+		// Ignore this slice if we don't own it.
+		if !h.Cluster.OwnsSlice(h.Host, slice) {
+			continue
+		}
+
+		// Otherwise retrieve the local fragment.
+		f, err := h.Index.CreateFragmentIfNotExists(db, frame, slice)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+
+		// Stream backup from remote node.
+		r, err := client.BackupSlice(db, frame, slice)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		} else if r == nil {
+			continue // slice doesn't exist
+		}
+
+		// Restore to local frame and always close reader.
+		if err := func() error {
+			defer r.Close()
+			if _, err := f.ReadFrom(r); err != nil {
+				return err
+			}
+			return nil
+		}(); err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+	}
+}
+
 // handleGetVersion handles /version requests.
 func (h *Handler) handleVersion(w http.ResponseWriter, r *http.Request) {
 	if err := json.NewEncoder(w).Encode(struct {