Merge pull request #1600 from tgruben/restore

[ CORE-525] Restore

api.go

@@ -17,6 +17,7 @@
 package pilosa
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/binary"
@@ -26,6 +27,7 @@ import (
 	"io/ioutil"
 	"math"
 	"net/url"
+	"os"
 	"sort"
 	"strconv"
 	"strings"
@@ -2200,6 +2202,9 @@ func (api *API) WriteIDAllocDataTo(w io.Writer) error {
 	_, err := api.holder.ida.WriteTo(w)
 	return err
 }
+func (api *API) RestoreIDAlloc(r io.Reader) error {
+	return api.holder.ida.Replace(r)
+}
 
 // TranslateIndexDB is an internal function to load the index keys database
 // rd is a boltdb file.
@@ -2219,6 +2224,92 @@ func (api *API) TranslateFieldDB(ctx context.Context, indexName, fieldName strin
 	return err
 }
 
+// RestoreShard
+func (api *API) RestoreShard(ctx context.Context, indexName string, shard uint64, rd io.Reader) error {
+	snap := topology.NewClusterSnapshot(api.cluster.noder, api.cluster.Hasher, api.cluster.ReplicaN)
+	if !snap.OwnsShard(api.server.nodeID, indexName, shard) {
+		return ErrClusterDoesNotOwnShard // TODO (twg)really just node doesn't own shard but leave for now
+	}
+
+	idx := api.holder.Index(indexName)
+	//need to get a dbShard
+	dbs, err := idx.Txf().dbPerShard.GetDBShard(indexName, shard, idx)
+	if err != nil {
+		return err
+	}
+	//need to find the path to the db
+	//will not work on blue green
+	db := dbs.W[0]
+	finalPath := db.Path() + "/data"
+	tempPath := finalPath + ".tmp"
+	o, err := os.OpenFile(tempPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
+	if err != nil {
+		return err
+	}
+	defer o.Close()
+
+	bw := bufio.NewWriter(o)
+	if _, err = io.Copy(bw, rd); err != nil {
+		return err
+	} else if err := bw.Flush(); err != nil {
+		return err
+	} else if err := o.Sync(); err != nil {
+		return err
+	} else if err := o.Close(); err != nil {
+		return err
+	}
+
+	if err != nil {
+		_ = os.Remove(tempPath)
+		return err
+	}
+	err = db.CloseDB()
+	if err != nil {
+		return err
+	}
+	err = os.Rename(tempPath, finalPath)
+	if err != nil {
+		_ = os.Remove(tempPath)
+		return err
+	}
+	err = db.OpenDB()
+	if err != nil {
+		return err
+	}
+	tx, err := db.NewTx(false, idx.name, Txo{})
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+	//arguments idx,shard do not matter for rbf they
+	//are ignored
+	flvs, err := tx.GetSortedFieldViewList(idx, shard)
+	if err != nil {
+		return nil
+	}
+
+	for _, flv := range flvs {
+		fld := idx.field(flv.Field)
+		view, ok := fld.viewMap[flv.View]
+		if !ok {
+			view, err = fld.createViewIfNotExists(flv.View)
+			if err != nil {
+				return err
+			}
+		}
+		frag, err := view.CreateFragmentIfNotExists(shard)
+		if err != nil {
+			return err
+		}
+		err = frag.RebuildRankCache(ctx)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 type serverInfo struct {
 	ShardWidth       uint64 `json:"shardWidth"`
 	ReplicaN         int    `json:"replicaN"`

bolt.go

@@ -217,6 +217,23 @@ func (w *BoltWrapper) HasData() (has bool, err error) {
 func (w *BoltWrapper) CleanupTx(tx Tx) {
 	// inlined into Rollback and Commit, so this is a no-op, just here to satisfy the interface.
 }
+func (w *BoltWrapper) CloseDB() error {
+	w.muDb.Lock()
+	defer w.muDb.Unlock()
+	w.closed = true
+	return w.db.Close()
+}
+func (w *BoltWrapper) OpenDB() error {
+	w.muDb.Lock()
+	defer w.muDb.Unlock()
+	db, err := bolt.Open(w.path, 0666, &bolt.Options{Timeout: 5 * time.Second, InitialMmapSize: TxInitialMmapSize})
+	if err != nil {
+		return err
+	}
+	w.db = db
+	w.closed = false
+	return nil
+}
 
 func (tx *BoltTx) IsDone() (done bool) {
 	return atomic.LoadInt64(&tx.unlocked) == 1

client.go

@@ -91,6 +91,9 @@ type InternalClient interface {
 
 	GetNodeUsage(ctx context.Context, uri *pnet.URI) (map[string]NodeUsage, error)
 	GetPastQueries(ctx context.Context, uri *pnet.URI) ([]PastQueryStatus, error)
+
+	ImportFieldKeys(ctx context.Context, uri *pnet.URI, index, field string, remote bool, rddbdata io.Reader) error
+	ImportIndexKeys(ctx context.Context, uri *pnet.URI, index string, partitionID int, remote bool, rddbdata io.Reader) error
 }
 
 //===============
@@ -274,3 +277,10 @@ func (n nopInternalClient) GetNodeUsage(ctx context.Context, uri *pnet.URI) (map
 func (n nopInternalClient) GetPastQueries(ctx context.Context, uri *pnet.URI) ([]PastQueryStatus, error) {
 	return nil, nil
 }
+func (c nopInternalClient) ImportFieldKeys(ctx context.Context, uri *pnet.URI, index, field string, remote bool, rddbdata io.Reader) error {
+	return nil
+}
+
+func (c nopInternalClient) ImportIndexKeys(ctx context.Context, uri *pnet.URI, index string, partitionID int, remote bool, rddbdata io.Reader) error {
+	return nil
+}

cmd/restore.go

@@ -0,0 +1,59 @@
+// Copyright 2020 Pilosa Corp.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"io"
+
+	"github.com/pilosa/pilosa/v2/ctl"
+	"github.com/spf13/cobra"
+)
+
+func newRestoreCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command {
+	c := ctl.NewRestoreCommand(stdin, stdout, stderr)
+	restoreCmd := &cobra.Command{
+		Use:   "restore [flags] PATH ",
+		Short: "restore a backup",
+		Long: `
+		The restore command will take a backup archive and restore it to a new, clean cluster.
+`,
+		Args: func(cmd *cobra.Command, args []string) error {
+			if len(args) == 0 {
+				return fmt.Errorf("data directory path required")
+			} else if len(args) > 1 {
+				return fmt.Errorf("too many command line arguments")
+			}
+			c.Path = args[0]
+			return nil
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return c.Run(context.Background())
+		},
+	}
+	flags := restoreCmd.Flags()
+	flags.StringVarP(&c.Host, "host", "", "localhost:10101", "host:port of Pilosa.")
+	flags.StringVarP(&c.Path, "source", "s", "", "pilosa backup file")
+	ctl.SetTLSConfig(
+		flags, "",
+		&c.TLS.CertificatePath,
+		&c.TLS.CertificateKeyPath,
+		&c.TLS.CACertPath,
+		&c.TLS.SkipVerify,
+		&c.TLS.EnableClientVerification)
+
+	return restoreCmd
+}

cmd/root.go

@@ -64,6 +64,7 @@ at https://www.pilosa.com/docs/.
 	rc.PersistentFlags().StringP("config", "c", "", "Configuration file to read from.")
 
 	rc.AddCommand(newBackupCommand(stdin, stdout, stderr))
+	rc.AddCommand(newRestoreCommand(stdin, stdout, stderr))
 	rc.AddCommand(newCheckCommand(stdin, stdout, stderr))
 	rc.AddCommand(newConfigCommand(stdin, stdout, stderr))
 	rc.AddCommand(newExportCommand(stdin, stdout, stderr))

ctl/restore.go

@@ -0,0 +1,192 @@
+// Copyright 2017 Pilosa Corp.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ctl
+
+import (
+	"archive/tar"
+	"compress/gzip"
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	gohttp "net/http"
+	"os"
+	"strconv"
+	"strings"
+
+	"github.com/pilosa/pilosa/v2"
+	"github.com/pilosa/pilosa/v2/server"
+	"github.com/pilosa/pilosa/v2/topology"
+)
+
+// RestoreCommand represents a command for restoring a backup to
+type RestoreCommand struct {
+	TLS  server.TLSConfig
+	Host string
+
+	// Filepath to the backup file.
+	Path string
+	// Reusable client.
+	client pilosa.InternalClient
+
+	// Standard input/output
+	*pilosa.CmdIO
+}
+
+// NewRestoreCommand returns a new instance of RestoreCommand.
+func NewRestoreCommand(stdin io.Reader, stdout, stderr io.Writer) *RestoreCommand {
+	return &RestoreCommand{
+		CmdIO: pilosa.NewCmdIO(stdin, stdout, stderr),
+	}
+}
+
+// Run executes the restore.
+func (cmd *RestoreCommand) Run(ctx context.Context) error {
+	logger := cmd.Logger()
+
+	// Validate arguments.
+	if cmd.Path == "" {
+		return fmt.Errorf("-s flag required")
+	}
+	useStdin := cmd.Path == "-"
+
+	var f *os.File
+	// read from Stdin if path specified as -
+	if useStdin {
+		f = os.Stdin
+	} else {
+		f, err := os.Open(cmd.Path)
+		if err != nil {
+			return (err)
+		}
+		defer f.Close()
+	}
+	// Create a client to the server.
+	client, err := commandClient(cmd)
+	if err != nil {
+		return fmt.Errorf("creating client: %w", err)
+	}
+	cmd.client = client
+	var tarReader *tar.Reader
+	if strings.HasSuffix(cmd.Path, "gz") {
+		gzf, err := gzip.NewReader(f)
+		if err != nil {
+			return err
+		}
+		tarReader = tar.NewReader(gzf)
+	} else {
+		tarReader = tar.NewReader(f)
+	}
+	nodes, err := cmd.client.Nodes(ctx)
+	if err != nil {
+		return err
+	}
+	var primary *topology.Node
+	for _, node := range nodes {
+		if node.IsPrimary {
+			primary = node
+			break
+		}
+
+	}
+	c := &gohttp.Client{}
+	if primary == nil {
+		return errors.New("no primary")
+	}
+	for {
+		header, err := tarReader.Next()
+		if err == io.EOF {
+			break
+		}
+		record := strings.Split(header.Name, "/")
+		if len(record) == 1 {
+			switch record[0] {
+			case "schema":
+				logger.Printf("Load Schema")
+				url := primary.URI.Path("/schema")
+				_, err = c.Post(url, "application/json", tarReader)
+				if err != nil {
+					return err
+				}
+			case "idalloc":
+				logger.Printf("Load ids")
+				url := primary.URI.Path("/internal/idalloc/restore")
+				_, err = c.Post(url, "application/octet-stream", tarReader)
+				if err != nil {
+					return err
+				}
+			default:
+				return err
+
+			}
+			continue
+		}
+		indexName := record[1]
+		switch record[2] {
+		case "shards":
+			shard, err := strconv.Atoi(record[3])
+			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)
+			if err != nil {
+				return err
+			}
+		case "translate":
+			partitionID, err := strconv.Atoi(record[3])
+			logger.Printf("column keys %v (%v)", indexName, partitionID)
+			if err != nil {
+				return err
+			}
+
+			err = cmd.client.ImportIndexKeys(ctx, &primary.URI, indexName, partitionID, false, tarReader)
+			if err != nil {
+				return err
+			}
+		case "attributes":
+			//skip
+		case "fields":
+			fieldName := record[3]
+			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)
+				if err != nil {
+					return err
+				}
+			case "attributes":
+			//skip
+			default:
+				return fmt.Errorf("unknown restore action: %v", action)
+			}
+
+		}
+
+	}
+	/*	Fetch the cluster nodes from the target host.
+		For each index:
+		Upload the RBF snapshot for each shard to the nodes that own the shard.
+		Upload the index & field translation BoltDB snapshots to each node.
+		If possible, trigger the node to reload itself. Otherwise a restart would be required.
+	*/
+
+	return nil
+}
+func (cmd *RestoreCommand) TLSHost() string { return cmd.Host }
+
+func (cmd *RestoreCommand) TLSConfiguration() server.TLSConfig { return cmd.TLS }

dbshard.go

@@ -66,6 +66,9 @@ type DBWrapper interface {
 	Path() string
 	HasData() (has bool, err error)
 	SetHolder(h *Holder)
+	//needed for restore
+	CloseDB() error
+	OpenDB() error
 }
 
 type DBRegistry interface {

executor.go

@@ -3688,7 +3688,6 @@ func (e *executor) executeRowsShard(ctx context.Context, qcx *Qcx, index string,
 		return nil, err
 	}
 	defer finisher(&err0)
-
 	for _, view := range views {
 		if err := ctx.Err(); err != nil {
 			return nil, err

field.go

@@ -1029,7 +1029,6 @@ func (f *Field) createViewIfNotExistsBase(cvm *CreateViewMessage) (*view, bool,
 			return nil, false, errors.Wrap(err, "persisting view")
 		}
 	}
-
 	view := f.newView(f.viewPath(cvm.View), cvm.View)
 
 	if err := view.openEmpty(); err != nil {

fragment.go

@@ -2890,6 +2890,31 @@ func (f *fragment) FlushCache() error {
 	defer f.mu.Unlock()
 	return f.flushCache()
 }
+func (f *fragment) RebuildRankCache(ctx context.Context) error {
+	if f.CacheType != CacheTypeRanked {
+		return nil //only rebuild ranked caches
+	}
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	tx, err := f.holder.BeginTx(false, f.idx, f.shard)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+	rows, err := f.unprotectedRows(ctx, tx, uint64(0))
+	if err != nil {
+		return err
+	}
+	for _, id := range rows {
+		n, err := tx.CountRange(f.index(), f.field(), f.view(), f.shard, id*ShardWidth, (id+1)*ShardWidth)
+		if err != nil {
+			return errors.Wrap(err, "CountRange")
+		}
+		f.cache.BulkAdd(id, n)
+	}
+	f.cache.Invalidate()
+	return nil
+}
 
 func (f *fragment) flushCache() error {
 	if f.cache == nil {
@@ -2909,6 +2934,9 @@ func (f *fragment) flushCache() error {
 		return errors.Wrap(err, "marshalling")
 	}
 
+	if err := os.MkdirAll(filepath.Dir(f.cachePath()), 0777); err != nil {
+		return errors.Wrap(err, "mkdir")
+	}
 	// Write to disk.
 	if err := ioutil.WriteFile(f.cachePath(), buf, 0666); err != nil {
 		return errors.Wrap(err, "writing")

http/client.go

@@ -199,7 +199,6 @@ func (c *InternalClient) Schema(ctx context.Context) ([]*pilosa.IndexInfo, error
 	}
 	return rsp.Indexes, nil
 }
-
 func (c *InternalClient) PostSchema(ctx context.Context, uri *pnet.URI, s *pilosa.Schema, remote bool) error {
 	u := uri.Path(fmt.Sprintf("/schema?remote=%v", remote))
 	buf, err := json.Marshal(s)

http/handler.go

@@ -438,9 +438,11 @@ func newRouter(handler *Handler) http.Handler {
 
 	router.HandleFunc("/internal/idalloc/reserve", handler.handleReserveIDs).Methods("POST").Name("ReserveIDs")
 	router.HandleFunc("/internal/idalloc/commit", handler.handleCommitIDs).Methods("POST").Name("CommitIDs")
+	router.HandleFunc("/internal/idalloc/restore", handler.handleRestoreIDAlloc).Methods("POST").Name("RestoreIDAllocData")
 	router.HandleFunc("/internal/idalloc/reset/{index}", handler.handleResetIDAlloc).Methods("POST").Name("ResetIDAlloc")
 	router.HandleFunc("/internal/idalloc/data", handler.handleIDAllocData).Methods("GET").Name("IDAllocData")
 
+	router.HandleFunc("/internal/restore/{index}/{shardID}", handler.handlePostRestore).Methods("POST").Name("Restore")
 	// endpoints for collecting cpu profiles from a chosen begin point to
 	// when the client wants to stop. Used for profiling imports that
 	// could be long or short.
@@ -2911,3 +2913,41 @@ func (h *Handler) handleIDAllocData(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 }
+
+func (h *Handler) handleRestoreIDAlloc(w http.ResponseWriter, r *http.Request) {
+	if err := h.api.RestoreIDAlloc(r.Body); err != nil {
+		http.Error(w, fmt.Sprintf("restoring id allocation: %v", err.Error()), http.StatusInternalServerError)
+		return
+	}
+	w.Header().Add("Content-Type", "text/plain")
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("OK")) //nolint:errcheck
+}
+func (h *Handler) handlePostRestore(w http.ResponseWriter, r *http.Request) {
+	indexName, ok := mux.Vars(r)["index"]
+	if !ok {
+		http.Error(w, "index name is required", http.StatusBadRequest)
+		return
+	}
+	shardID, ok := mux.Vars(r)["shardID"]
+	if !ok {
+		http.Error(w, "shardID is required", http.StatusBadRequest)
+		return
+	}
+	shard, err := strconv.ParseUint(shardID, 10, 64)
+	if err != nil {
+		http.Error(w, fmt.Sprintf("failed to parse shard %v %v err:%v", indexName, shardID, err), http.StatusBadRequest)
+		return
+	}
+	ctx := context.Background()
+	//validate shard for this node
+	err = h.api.RestoreShard(ctx, indexName, shard, r.Body)
+	if err != nil {
+		http.Error(w, fmt.Sprintf("failed to restore shared %v %v err:%v", indexName, shard, err), http.StatusBadRequest)
+		return
+	}
+
+	w.Header().Add("Content-Type", "text/plain")
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("OK")) //nolint:errcheck
+}

idalloc.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"io"
 	"math/bits"
+	"os"
 	"sort"
 	"time"
 
@@ -64,6 +65,37 @@ func OpenIDAllocator(path string) (*idAllocator, error) {
 	}
 	return &idAllocator{db}, nil
 }
+func (ida *idAllocator) Replace(reader io.Reader) error {
+	newFile := ida.db.Path() + ".bak"
+	liveFile := ida.db.Path()
+	file, err := os.OpenFile(newFile, os.O_RDWR|os.O_CREATE, 0600)
+	if err != nil {
+		return nil
+	}
+	_, err = io.Copy(file, reader)
+	if err != nil {
+		return err
+	}
+	file.Close()
+	err = ida.db.Close()
+	if err != nil {
+		return err
+	}
+	err = os.Rename(liveFile, liveFile+".sav")
+	if err != nil {
+		return err
+	}
+	err = os.Rename(newFile, liveFile)
+	if err != nil {
+		err = os.Rename(liveFile+".sav", liveFile)
+		return err
+	} else {
+		_ = os.Remove(liveFile + ".sav")
+	}
+	db, err := bolt.Open(liveFile, 0666, &bolt.Options{Timeout: 1 * time.Second})
+	ida.db = db
+	return err
+}
 
 func (ida *idAllocator) Close() error {
 	if ida == nil || ida.db == nil {

rbf.go

@@ -571,6 +571,26 @@ func (w *RbfDBWrapper) Close() error {
 	return w.db.Close()
 }
 
+// needed to handle the special case on reload, the close method unregisters the wrapper and all that is
+// required is the backing file get reloaded
+
+func (w *RbfDBWrapper) CloseDB() error {
+	w.muDb.Lock()
+	defer w.muDb.Unlock()
+	w.closed = true
+	return w.db.Close()
+}
+func (w *RbfDBWrapper) OpenDB() error {
+	w.muDb.Lock()
+	defer w.muDb.Unlock()
+	err := w.db.Open()
+	if err != nil {
+		return err
+	}
+	w.closed = false
+	return nil
+}
+
 var globalNextTxSnRBFTx int64
 
 func (w *RbfDBWrapper) NewTx(write bool, initialIndex string, o Txo) (_ Tx, err error) {

rrtx.go

@@ -663,6 +663,13 @@ func (w *RoaringWrapper) OpenSnList() (slc []int64) {
 	return nil
 }
 
+func (w *RoaringWrapper) CloseDB() error {
+	return errors.New("CloseDB not supported in roaring")
+}
+func (w *RoaringWrapper) OpenDB() error {
+	return errors.New("OpenDB not supported in roaring")
+}
+
 // statically confirm that RoaringTx satisfies the Tx interface.
 var _ Tx = (*RoaringTx)(nil)
 

server.go

@@ -493,7 +493,6 @@ func NewServer(opts ...ServerOption) (*Server, error) {
 	s.holder.schemator = s.schemator
 	s.holder.sharder = s.sharder
 	s.holder.serializer = s.serializer
-
 	return s, nil
 }