restore idalloc

api.go

@@ -40,7 +40,6 @@ import (
 	"github.com/pilosa/pilosa/v2/stats"
 	"github.com/pilosa/pilosa/v2/topology"
 	"github.com/pilosa/pilosa/v2/tracing"
-	"github.com/pilosa/pilosa/v2/vprint"
 	"github.com/pkg/errors"
 	"golang.org/x/sync/errgroup"
 )
@@ -2203,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.
@@ -2240,16 +2242,13 @@ func (api *API) RestoreShard(ctx context.Context, indexName string, shard uint64
 	db := dbs.W[0]
 	finalPath := db.Path() + "/data"
 	tempPath := finalPath + ".tmp"
-	vprint.VV("restore to %v", tempPath)
 	o, err := os.OpenFile(tempPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
 	if err != nil {
 		return err
 	}
 	w := bufio.NewWriter(o)
 	//close db if open
-	vprint.VV("restore index:%v shard:%v", indexName, shard)
-	n, err := io.Copy(w, rd)
-	vprint.VV("Written:%v", n)
+	_, err = io.Copy(w, rd)
 	w.Flush()
 	o.Close()
 	if err != nil {
@@ -2260,7 +2259,6 @@ func (api *API) RestoreShard(ctx context.Context, indexName string, shard uint64
 	if err != nil {
 		return err
 	}
-	vprint.VV("Rename %v to %v", tempPath, finalPath)
 	err = os.Rename(tempPath, finalPath)
 	if err != nil {
 		_ = os.Remove(tempPath)

ctl/restore.go

@@ -155,7 +155,6 @@ func (cmd *RestoreCommand) Run(ctx context.Context) error {
 				vprint.VV("Load Schema")
 				url := primary.URI.Path("/schema")
 				vprint.VV("SCHEMA %v", url)
-				//schemaBytes, err := ioutil.ReadAll(tarReader)
 				_, err = c.Post(url, "application/json", tarReader)
 				if err != nil {
 					return err
@@ -163,6 +162,11 @@ func (cmd *RestoreCommand) Run(ctx context.Context) error {
 
 			case "idalloc":
 				vprint.VV("Load ids")
+				url := primary.URI.Path("/internal/idalloc/restore")
+				_, err = c.Post(url, "application/octet-stream", tarReader)
+				if err != nil {
+					return err
+				}
 			default:
 				panic("UNKNOWN " + record[0])
 
@@ -196,7 +200,8 @@ func (cmd *RestoreCommand) Run(ctx context.Context) error {
 				return err
 			}
 		case "attributes":
-			vprint.VV("column attributes %v", indexName)
+			//skip
+			//vprint.VV("column attributes %v", indexName)
 		case "fields":
 			fieldName := record[3]
 			switch action := record[4]; action {
@@ -207,7 +212,7 @@ func (cmd *RestoreCommand) Run(ctx context.Context) error {
 					return err
 				}
 			case "attributes":
-				vprint.VV("field attributes %v %v", indexName, fieldName)
+			//	vprint.VV("field attributes %v %v", indexName, fieldName)
 			default:
 				panic("unknown:" + action)
 			}

http/handler.go

@@ -438,6 +438,7 @@ 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")
 
@@ -2913,6 +2914,15 @@ func (h *Handler) handleIDAllocData(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+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) {
 	/*
 		if !validHeaderAcceptType(r.Header, "text", "plain") {

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 {