mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-08-28 02:44:59 +00:00
limit memory for backuptar/restoretar (#2270)
This commit is contained in:
tgruben
GitHub
parent
528ebc93db
commit
c294bc70dc
5 changed files with 222 additions and 130 deletions
110
buffer/filebuffer.go
Normal file
110
buffer/filebuffer.go
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
package buffer
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"os"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// NewFileBuffer returns a file buffer which will use an in-memory buffer, until `max` bytes have been written, at which point it will write the contents of memory to a file, and continue writing future data to the file.
|
||||
// The file will be written to `temp` directory. The buffer fulfills the io.Reader and io.Writer interface
|
||||
func NewFileBuffer(max int, temp string) *FileBuffer {
|
||||
return &FileBuffer{max: max, tempDir: temp}
|
||||
}
|
||||
|
||||
type FileBuffer struct {
|
||||
max int
|
||||
buf bytes.Buffer
|
||||
file *os.File
|
||||
tempDir string
|
||||
reading bool
|
||||
files []*os.File
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func (fb *FileBuffer) Write(p []byte) (n int, err error) {
|
||||
if fb.reading {
|
||||
panic("cannot write after read")
|
||||
}
|
||||
if fb.file != nil {
|
||||
return fb.file.Write(p)
|
||||
}
|
||||
n, err = fb.buf.Write(p)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if fb.buf.Len() > fb.max {
|
||||
fb.file, err = ioutil.TempFile(fb.tempDir, "filebuffer-")
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_, err = io.Copy(fb.file, &fb.buf)
|
||||
fb.buf.Reset()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (fb *FileBuffer) Len() (int64, error) {
|
||||
if fb.file == nil {
|
||||
return int64(fb.buf.Len()), nil
|
||||
}
|
||||
fi, err := fb.file.Stat()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
return fi.Size(), nil
|
||||
}
|
||||
|
||||
func (fb *FileBuffer) Read(p []byte) (n int, err error) {
|
||||
if fb.file != nil {
|
||||
if !fb.reading {
|
||||
fb.reading = true
|
||||
_, err = fb.file.Seek(0, 0)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
return fb.file.Read(p)
|
||||
}
|
||||
fb.reading = true
|
||||
return fb.buf.Read(p)
|
||||
}
|
||||
|
||||
func (fb *FileBuffer) Close() error {
|
||||
if fb.file != nil {
|
||||
name := fb.file.Name()
|
||||
if err := fb.file.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, f := range fb.files {
|
||||
f.Close()
|
||||
}
|
||||
fb.files = fb.files[:0]
|
||||
fb.file = nil
|
||||
return os.Remove(name)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (fb *FileBuffer) Reset() error {
|
||||
fb.mu.Lock()
|
||||
defer fb.mu.Unlock()
|
||||
fb.reading = false
|
||||
fb.buf.Reset()
|
||||
return fb.Close()
|
||||
}
|
||||
|
||||
func (fb *FileBuffer) NewReader() (io.Reader, error) {
|
||||
fb.mu.Lock()
|
||||
defer fb.mu.Unlock()
|
||||
fb.reading = true
|
||||
if fb.file == nil {
|
||||
return bytes.NewReader(fb.buf.Bytes()), nil
|
||||
}
|
||||
f, err := os.OpenFile(fb.file.Name(), os.O_RDONLY, 0)
|
||||
fb.files = append(fb.files, f)
|
||||
return f, err
|
||||
}
|
||||
|
|
@ -28,6 +28,7 @@ Backs up a FeatureBase server to a local, tar-formatted snapshot file.
|
|||
ctl.SetTLSConfig(flags, "", &cmd.TLS.CertificatePath, &cmd.TLS.CertificateKeyPath, &cmd.TLS.CACertPath, &cmd.TLS.SkipVerify, &cmd.TLS.EnableClientVerification)
|
||||
flags.StringVar(&cmd.AuthToken, "auth-token", "", "Authentication token")
|
||||
flags.StringVar(&cmd.HeaderTimeoutStr, "header-timeout", cmd.HeaderTimeoutStr, "Length of time to wait for initial HTTP response before giving up.")
|
||||
flags.StringVar(&cmd.TempDir, "temp-dir", cmd.TempDir, "Location of temporary spillover files. The default is the system's default (usually /tmp)")
|
||||
|
||||
return ccmd
|
||||
}
|
||||
|
|
|
|||
|
|
@ -23,6 +23,8 @@ The Restore command will take a tar-formatted backup archive and restore it to a
|
|||
flags.DurationVar(&cmd.RetryPeriod, "retry-period", cmd.RetryPeriod, "Length of time after HTTP request failure to continue retrying request.")
|
||||
flags.StringVar(&cmd.Pprof, "pprof", cmd.Pprof, "host:port to listen for profiling requests at /debug/pprof and /debug/fgprof.")
|
||||
flags.StringVar(&cmd.AuthToken, "auth-token", "", "Authentication token")
|
||||
flags.StringVar(&cmd.TempDir, "temp-dir", cmd.TempDir, "Location of temporary spillover files. The default is the system's default (usually /tmp)")
|
||||
|
||||
ctl.SetTLSConfig(
|
||||
flags, "",
|
||||
&cmd.TLS.CertificatePath,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ package ctl
|
|||
|
||||
import (
|
||||
"archive/tar"
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/json"
|
||||
|
|
@ -17,6 +16,7 @@ import (
|
|||
|
||||
pilosa "github.com/featurebasedb/featurebase/v3"
|
||||
"github.com/featurebasedb/featurebase/v3/authn"
|
||||
"github.com/featurebasedb/featurebase/v3/buffer"
|
||||
"github.com/featurebasedb/featurebase/v3/disco"
|
||||
"github.com/featurebasedb/featurebase/v3/encoding/proto"
|
||||
"github.com/featurebasedb/featurebase/v3/logger"
|
||||
|
|
@ -37,6 +37,9 @@ type BackupTarCommand struct { // nolint: maligned
|
|||
// Path to write the backup to.
|
||||
OutputPath string
|
||||
|
||||
// TempDir location of scratch files
|
||||
TempDir string
|
||||
|
||||
// Amount of time after first failed request to continue retrying.
|
||||
RetryPeriod time.Duration `json:"retry-period"`
|
||||
|
||||
|
|
@ -164,17 +167,16 @@ func (cmd *BackupTarCommand) Run(ctx context.Context) (err error) {
|
|||
tw := tar.NewWriter(w)
|
||||
defer tw.Close()
|
||||
|
||||
buf := new(bytes.Buffer)
|
||||
// Backup schema.
|
||||
if err := cmd.backupTarSchema(ctx, tw, schema); err != nil {
|
||||
return fmt.Errorf("cannot back up schema: %w", err)
|
||||
} else if err := cmd.backupTarIDAllocData(ctx, tw, buf); err != nil {
|
||||
} else if err := cmd.backupTarIDAllocData(ctx, tw); err != nil {
|
||||
return fmt.Errorf("cannot back up id alloc data: %w", err)
|
||||
}
|
||||
|
||||
// Backup data for each index.
|
||||
for _, ii := range schema.Indexes {
|
||||
if err := cmd.backupTarIndex(ctx, tw, ii, buf); err != nil {
|
||||
if err := cmd.backupTarIndex(ctx, tw, ii); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
|
@ -220,8 +222,7 @@ func (cmd *BackupTarCommand) backupTarSchema(ctx context.Context, tw *tar.Writer
|
|||
return nil
|
||||
}
|
||||
|
||||
func (cmd *BackupTarCommand) backupTarIDAllocData(ctx context.Context, tw *tar.Writer, buf *bytes.Buffer) error {
|
||||
buf.Reset()
|
||||
func (cmd *BackupTarCommand) backupTarIDAllocData(ctx context.Context, tw *tar.Writer) error {
|
||||
logger := cmd.Logger()
|
||||
logger.Printf("backing up id alloc data")
|
||||
|
||||
|
|
@ -231,28 +232,11 @@ func (cmd *BackupTarCommand) backupTarIDAllocData(ctx context.Context, tw *tar.W
|
|||
}
|
||||
defer rc.Close()
|
||||
|
||||
// Read to buffer to determine size.
|
||||
if _, err := buf.ReadFrom(rc); err != nil {
|
||||
return fmt.Errorf("copying id alloc data to memory: %w", err)
|
||||
}
|
||||
|
||||
// Build header & copy data to archive.
|
||||
if err = tw.WriteHeader(&tar.Header{
|
||||
Name: "idalloc",
|
||||
Mode: 0o666,
|
||||
Size: int64(buf.Len()),
|
||||
ModTime: time.Now(),
|
||||
}); err != nil {
|
||||
return err
|
||||
} else if _, err := io.Copy(tw, buf); err != nil {
|
||||
return fmt.Errorf("copying id alloc data to archive: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
return writeToTar(tw, "idalloc", rc, cmd.TempDir)
|
||||
}
|
||||
|
||||
// backupTarIndex backs up all shards for a given index.
|
||||
func (cmd *BackupTarCommand) backupTarIndex(ctx context.Context, tw *tar.Writer, ii *pilosa.IndexInfo, buf *bytes.Buffer) error {
|
||||
func (cmd *BackupTarCommand) backupTarIndex(ctx context.Context, tw *tar.Writer, ii *pilosa.IndexInfo) error {
|
||||
logger := cmd.Logger()
|
||||
logger.Printf("backing up index: %q", ii.Name)
|
||||
|
||||
|
|
@ -263,14 +247,14 @@ func (cmd *BackupTarCommand) backupTarIndex(ctx context.Context, tw *tar.Writer,
|
|||
|
||||
// Back up all bitmap data for the index.
|
||||
for _, shard := range shards {
|
||||
if err := cmd.backupTarShard(ctx, tw, ii.Name, shard, buf); err != nil {
|
||||
if err := cmd.backupTarShard(ctx, tw, ii.Name, shard); err != nil {
|
||||
return fmt.Errorf("cannot backup shard %d on index %q: %w", shard, ii.Name, err)
|
||||
}
|
||||
}
|
||||
|
||||
if ii.Options.Keys {
|
||||
// Back up translation data after bitmap data so we ensurean translate all data.
|
||||
if err := cmd.backupTarIndexTranslateData(ctx, tw, ii.Name, buf); err != nil {
|
||||
if err := cmd.backupTarIndexTranslateData(ctx, tw, ii.Name); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
|
@ -280,7 +264,7 @@ func (cmd *BackupTarCommand) backupTarIndex(ctx context.Context, tw *tar.Writer,
|
|||
if !fi.Options.Keys {
|
||||
continue
|
||||
}
|
||||
if err := cmd.backupTarFieldTranslateData(ctx, tw, ii.Name, fi.Name, buf); err != nil {
|
||||
if err := cmd.backupTarFieldTranslateData(ctx, tw, ii.Name, fi.Name); err != nil {
|
||||
return fmt.Errorf("cannot backup field translation data for field %q on index %q: %w", fi.Name, ii.Name, err)
|
||||
}
|
||||
}
|
||||
|
|
@ -289,7 +273,7 @@ func (cmd *BackupTarCommand) backupTarIndex(ctx context.Context, tw *tar.Writer,
|
|||
}
|
||||
|
||||
// backupTarShard backs up a single shard from a single index.
|
||||
func (cmd *BackupTarCommand) backupTarShard(ctx context.Context, tw *tar.Writer, indexName string, shard uint64, buf *bytes.Buffer) (err error) {
|
||||
func (cmd *BackupTarCommand) backupTarShard(ctx context.Context, tw *tar.Writer, indexName string, shard uint64) (err error) {
|
||||
nodes, err := cmd.client.FragmentNodes(ctx, indexName, shard)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot determine fragment nodes: %w", err)
|
||||
|
|
@ -298,7 +282,7 @@ func (cmd *BackupTarCommand) backupTarShard(ctx context.Context, tw *tar.Writer,
|
|||
}
|
||||
|
||||
for _, node := range nodes {
|
||||
if e := cmd.backupTarShardNode(ctx, tw, indexName, shard, node, buf); e == nil {
|
||||
if e := cmd.backupTarShardNode(ctx, tw, indexName, shard, node); e == nil {
|
||||
break
|
||||
} else if err == nil {
|
||||
err = e // save first error, try next node
|
||||
|
|
@ -306,7 +290,7 @@ func (cmd *BackupTarCommand) backupTarShard(ctx context.Context, tw *tar.Writer,
|
|||
}
|
||||
|
||||
for _, node := range nodes {
|
||||
if e := cmd.backupTarShardDataframe(ctx, tw, indexName, shard, node, buf); e == nil {
|
||||
if e := cmd.backupTarShardDataframe(ctx, tw, indexName, shard, node); e == nil {
|
||||
break
|
||||
} else if err == nil {
|
||||
err = e // save first error, try next node
|
||||
|
|
@ -316,8 +300,7 @@ func (cmd *BackupTarCommand) backupTarShard(ctx context.Context, tw *tar.Writer,
|
|||
}
|
||||
|
||||
// backupTarShardNode backs up a single shard from a single index on a specific node.
|
||||
func (cmd *BackupTarCommand) backupTarShardNode(ctx context.Context, tw *tar.Writer, indexName string, shard uint64, node *disco.Node, buf *bytes.Buffer) error {
|
||||
buf.Reset()
|
||||
func (cmd *BackupTarCommand) backupTarShardNode(ctx context.Context, tw *tar.Writer, indexName string, shard uint64, node *disco.Node) error {
|
||||
logger := cmd.Logger()
|
||||
logger.Printf("backing up shard: index=%q id=%d", indexName, shard)
|
||||
|
||||
|
|
@ -332,30 +315,10 @@ func (cmd *BackupTarCommand) backupTarShardNode(ctx context.Context, tw *tar.Wri
|
|||
return fmt.Errorf("fetching shard reader: %w", err)
|
||||
}
|
||||
defer rc.Close()
|
||||
|
||||
// Read to buffer to determine size.
|
||||
// TODO: Provide size via the reader itself.
|
||||
if _, err := buf.ReadFrom(rc); err != nil {
|
||||
return fmt.Errorf("copying shard data to memory: %w", err)
|
||||
}
|
||||
|
||||
// Build header & copy data to archive.
|
||||
if err = tw.WriteHeader(&tar.Header{
|
||||
Name: filename,
|
||||
Mode: 0o666,
|
||||
Size: int64(buf.Len()),
|
||||
ModTime: time.Now(),
|
||||
}); err != nil {
|
||||
return err
|
||||
} else if _, err := io.Copy(tw, buf); err != nil {
|
||||
return fmt.Errorf("copying shard data to archive: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
return writeToTar(tw, filename, rc, cmd.TempDir)
|
||||
}
|
||||
|
||||
func (cmd *BackupTarCommand) backupTarShardDataframe(ctx context.Context, tw *tar.Writer, indexName string, shard uint64, node *disco.Node, buf *bytes.Buffer) error {
|
||||
buf.Reset()
|
||||
func (cmd *BackupTarCommand) backupTarShardDataframe(ctx context.Context, tw *tar.Writer, indexName string, shard uint64, node *disco.Node) error {
|
||||
logger := cmd.Logger()
|
||||
logger.Printf("backing up dataframe shard: index=%q shard=%d", indexName, shard)
|
||||
|
||||
|
|
@ -375,38 +338,21 @@ func (cmd *BackupTarCommand) backupTarShardDataframe(ctx context.Context, tw *ta
|
|||
}
|
||||
|
||||
filename := filepath.Join("indexes", indexName, "dataframe", fmt.Sprintf("%04d", shard))
|
||||
logger.Printf("writing %v", filename)
|
||||
if _, err := buf.ReadFrom(resp.Body); err != nil {
|
||||
return fmt.Errorf("copying shard data to memory: %w", err)
|
||||
}
|
||||
|
||||
// Build header & copy data to archive.
|
||||
if err = tw.WriteHeader(&tar.Header{
|
||||
Name: filename,
|
||||
Mode: 0o666,
|
||||
Size: int64(buf.Len()),
|
||||
ModTime: time.Now(),
|
||||
}); err != nil {
|
||||
return err
|
||||
} else if _, err := io.Copy(tw, buf); err != nil {
|
||||
return fmt.Errorf("copying shard data to archive: %w", err)
|
||||
}
|
||||
return nil
|
||||
return writeToTar(tw, filename, resp.Body, cmd.TempDir)
|
||||
}
|
||||
|
||||
func (cmd *BackupTarCommand) backupTarIndexTranslateData(ctx context.Context, tw *tar.Writer, name string, buf *bytes.Buffer) error {
|
||||
func (cmd *BackupTarCommand) backupTarIndexTranslateData(ctx context.Context, tw *tar.Writer, name string) error {
|
||||
// TODO: Fetch holder partition count.
|
||||
partitionN := disco.DefaultPartitionN
|
||||
for partitionID := 0; partitionID < partitionN; partitionID++ {
|
||||
if err := cmd.backupTarIndexPartitionTranslateData(ctx, tw, name, partitionID, buf); err != nil {
|
||||
if err := cmd.backupTarIndexPartitionTranslateData(ctx, tw, name, partitionID); err != nil {
|
||||
return fmt.Errorf("cannot backup index translation data for partition %d on %q: %w", partitionID, name, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (cmd *BackupTarCommand) backupTarIndexPartitionTranslateData(ctx context.Context, tw *tar.Writer, name string, partitionID int, buf *bytes.Buffer) error {
|
||||
buf.Reset()
|
||||
func (cmd *BackupTarCommand) backupTarIndexPartitionTranslateData(ctx context.Context, tw *tar.Writer, name string, partitionID int) error {
|
||||
logger := cmd.Logger()
|
||||
logger.Printf("backing up index translation data: %s/%d", name, partitionID)
|
||||
|
||||
|
|
@ -418,28 +364,10 @@ func (cmd *BackupTarCommand) backupTarIndexPartitionTranslateData(ctx context.Co
|
|||
}
|
||||
defer rc.Close()
|
||||
|
||||
// Read to buffer to determine size.
|
||||
if _, err := buf.ReadFrom(rc); err != nil {
|
||||
return fmt.Errorf("copying translate data to memory: %w", err)
|
||||
}
|
||||
|
||||
// Build header & copy data to archive.
|
||||
if err = tw.WriteHeader(&tar.Header{
|
||||
Name: path.Join("indexes", name, "translate", fmt.Sprintf("%04d", partitionID)),
|
||||
Mode: 0o666,
|
||||
Size: int64(buf.Len()),
|
||||
ModTime: time.Now(),
|
||||
}); err != nil {
|
||||
return err
|
||||
} else if _, err := io.Copy(tw, buf); err != nil {
|
||||
return fmt.Errorf("copying translate data to archive: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
return writeToTar(tw, path.Join("indexes", name, "translate", fmt.Sprintf("%04d", partitionID)), rc, cmd.TempDir)
|
||||
}
|
||||
|
||||
func (cmd *BackupTarCommand) backupTarFieldTranslateData(ctx context.Context, tw *tar.Writer, indexName, fieldName string, buf *bytes.Buffer) error {
|
||||
buf.Reset()
|
||||
func (cmd *BackupTarCommand) backupTarFieldTranslateData(ctx context.Context, tw *tar.Writer, indexName, fieldName string) error {
|
||||
logger := cmd.Logger()
|
||||
logger.Printf("backing up field translation data: %s/%s", indexName, fieldName)
|
||||
|
||||
|
|
@ -450,17 +378,37 @@ func (cmd *BackupTarCommand) backupTarFieldTranslateData(ctx context.Context, tw
|
|||
return fmt.Errorf("fetching translate data reader: %w", err)
|
||||
}
|
||||
defer rc.Close()
|
||||
return writeToTar(tw, path.Join("indexes", indexName, "fields", fieldName, "translate"), rc, cmd.TempDir)
|
||||
}
|
||||
|
||||
func (cmd *BackupTarCommand) TLSHost() string { return cmd.Host }
|
||||
|
||||
func (cmd *BackupTarCommand) TLSConfiguration() server.TLSConfig { return cmd.TLS }
|
||||
|
||||
func writeToTar(tw *tar.Writer, entryName string, rc io.Reader, tmpDir string) error {
|
||||
spillFile, err := os.CreateTemp("", "spill")
|
||||
if err != nil {
|
||||
return fmt.Errorf("creating temp file : %w", err)
|
||||
}
|
||||
defer func() {
|
||||
spillFile.Close()
|
||||
os.Remove(spillFile.Name())
|
||||
}()
|
||||
mb512 := 2 << 29
|
||||
buf := buffer.NewFileBuffer(mb512, tmpDir)
|
||||
defer buf.Close()
|
||||
|
||||
n, err := io.Copy(buf, rc)
|
||||
// Read to buffer to determine size.
|
||||
if _, err := buf.ReadFrom(rc); err != nil {
|
||||
if err != nil {
|
||||
return fmt.Errorf("copying translate data to memory: %w", err)
|
||||
}
|
||||
|
||||
// Build header & copy data to archive.
|
||||
if err = tw.WriteHeader(&tar.Header{
|
||||
Name: path.Join("indexes", indexName, "fields", fieldName, "translate"),
|
||||
Name: entryName,
|
||||
Mode: 0o666,
|
||||
Size: int64(buf.Len()),
|
||||
Size: n,
|
||||
ModTime: time.Now(),
|
||||
}); err != nil {
|
||||
return err
|
||||
|
|
@ -469,7 +417,3 @@ func (cmd *BackupTarCommand) backupTarFieldTranslateData(ctx context.Context, tw
|
|||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (cmd *BackupTarCommand) TLSHost() string { return cmd.Host }
|
||||
|
||||
func (cmd *BackupTarCommand) TLSConfiguration() server.TLSConfig { return cmd.TLS }
|
||||
|
|
|
|||
|
|
@ -3,12 +3,13 @@ package ctl
|
|||
|
||||
import (
|
||||
"archive/tar"
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
gohttp "net/http"
|
||||
"os"
|
||||
"strconv"
|
||||
|
|
@ -17,6 +18,7 @@ import (
|
|||
|
||||
pilosa "github.com/featurebasedb/featurebase/v3"
|
||||
"github.com/featurebasedb/featurebase/v3/authn"
|
||||
"github.com/featurebasedb/featurebase/v3/buffer"
|
||||
"github.com/featurebasedb/featurebase/v3/disco"
|
||||
"github.com/featurebasedb/featurebase/v3/logger"
|
||||
"github.com/featurebasedb/featurebase/v3/server"
|
||||
|
|
@ -48,6 +50,9 @@ type RestoreTarCommand struct {
|
|||
TLS server.TLSConfig
|
||||
|
||||
AuthToken string
|
||||
|
||||
// TempDir location of scratch files
|
||||
TempDir string
|
||||
}
|
||||
|
||||
// Logger returns the command's associated Logger to maintain CommandWithTLSSupport interface compatibility
|
||||
|
|
@ -134,7 +139,10 @@ func (cmd *RestoreTarCommand) Run(ctx context.Context) (err error) {
|
|||
return errors.New("no primary")
|
||||
}
|
||||
c := &gohttp.Client{}
|
||||
buf := new(bytes.Buffer)
|
||||
// buf := new(bytes.Buffer)
|
||||
mb512 := 2 << 29
|
||||
buf := buffer.NewFileBuffer(mb512, cmd.TempDir)
|
||||
defer buf.Reset()
|
||||
for {
|
||||
buf.Reset()
|
||||
header, err := tarReader.Next()
|
||||
|
|
@ -177,22 +185,21 @@ func (cmd *RestoreTarCommand) Run(ctx context.Context) (err error) {
|
|||
} else if len(fragmentNodes) == 0 {
|
||||
return fmt.Errorf("no fragmentNodes available")
|
||||
}
|
||||
|
||||
_, err = io.Copy(buf, tarReader)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "copying")
|
||||
}
|
||||
|
||||
g, _ := errgroup.WithContext(ctx)
|
||||
for _, node := range fragmentNodes {
|
||||
node := node
|
||||
rd, err := buf.NewReader()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
g.Go(func() error {
|
||||
client := &gohttp.Client{}
|
||||
rd := bytes.NewReader(buf.Bytes())
|
||||
logger.Printf("shard %v %v", shard, indexName)
|
||||
url := node.URI.Path(fmt.Sprintf("/internal/restore/%v/%v", indexName, shard))
|
||||
_, err = client.Post(url, "application/octet-stream", rd)
|
||||
return err
|
||||
return Post(ctx, url, "application/octet-stream", rd, nil)
|
||||
})
|
||||
}
|
||||
if err := g.Wait(); err != nil {
|
||||
|
|
@ -218,13 +225,14 @@ func (cmd *RestoreTarCommand) Run(ctx context.Context) (err error) {
|
|||
g, _ := errgroup.WithContext(ctx)
|
||||
for _, node := range fragmentNodes {
|
||||
node := node
|
||||
rd, err := buf.NewReader()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
g.Go(func() error {
|
||||
client := &gohttp.Client{}
|
||||
rd := bytes.NewReader(buf.Bytes())
|
||||
logger.Printf("dataframe shard %v %v", shard, indexName)
|
||||
url := node.URI.Path(fmt.Sprintf("/internal/dataframe/restore/%v/%v", indexName, shard))
|
||||
_, err = client.Post(url, "application/octet-stream", rd)
|
||||
return err
|
||||
return Post(ctx, url, "application/octet-stream", rd, nil)
|
||||
})
|
||||
}
|
||||
if err := g.Wait(); err != nil {
|
||||
|
|
@ -249,13 +257,13 @@ func (cmd *RestoreTarCommand) Run(ctx context.Context) (err error) {
|
|||
g, _ := errgroup.WithContext(ctx)
|
||||
for _, node := range partitionNodes {
|
||||
node := node
|
||||
rd, err := buf.NewReader()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
g.Go(func() error {
|
||||
// rd := bytes.NewReader(shardBytes)
|
||||
rd := func() (io.Reader, error) {
|
||||
return bytes.NewReader(buf.Bytes()), nil
|
||||
}
|
||||
|
||||
return cmd.client.ImportIndexKeys(ctx, &node.URI, indexName, partitionID, false, rd)
|
||||
url := node.URI.Path(fmt.Sprintf("/internal/translate/index/%s/%d", indexName, partitionID))
|
||||
return Post(ctx, url, "application/octet-stream", rd, nil)
|
||||
})
|
||||
}
|
||||
if err := g.Wait(); err != nil {
|
||||
|
|
@ -269,23 +277,20 @@ func (cmd *RestoreTarCommand) Run(ctx context.Context) (err error) {
|
|||
switch action := record[4]; action {
|
||||
case "translate":
|
||||
logger.Printf("field keys %v %v", indexName, fieldName)
|
||||
// needs to go to all nodes
|
||||
|
||||
_, err = io.Copy(buf, tarReader)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "copying")
|
||||
}
|
||||
|
||||
g, _ := errgroup.WithContext(ctx)
|
||||
for _, node := range nodes {
|
||||
node := node
|
||||
rd, err := buf.NewReader()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
g.Go(func() error {
|
||||
// rd := bytes.NewReader(shardBytes)
|
||||
rd := func() (io.Reader, error) {
|
||||
return bytes.NewReader(buf.Bytes()), nil
|
||||
}
|
||||
|
||||
return cmd.client.ImportFieldKeys(ctx, &node.URI, indexName, fieldName, false, rd)
|
||||
url := node.URI.Path(fmt.Sprintf("/internal/translate/field/%s/%s", indexName, fieldName))
|
||||
return Post(ctx, url, "application/octet-stream", rd, nil)
|
||||
})
|
||||
}
|
||||
if err := g.Wait(); err != nil {
|
||||
|
|
@ -312,3 +317,33 @@ func (cmd *RestoreTarCommand) Run(ctx context.Context) (err error) {
|
|||
func (cmd *RestoreTarCommand) TLSHost() string { return cmd.Host }
|
||||
|
||||
func (cmd *RestoreTarCommand) TLSConfiguration() server.TLSConfig { return cmd.TLS }
|
||||
|
||||
func Post(ctx context.Context, url, contentType string, rd io.Reader, query map[string]string) error {
|
||||
client := &gohttp.Client{}
|
||||
req, err := http.NewRequest(http.MethodPost, url, rd)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("User-Agent", "pilosa/"+pilosa.Version)
|
||||
req.Header.Set("Content-Type", contentType)
|
||||
pilosa.AddAuthToken(ctx, &req.Header)
|
||||
|
||||
// appending to existing query args
|
||||
|
||||
q := req.URL.Query()
|
||||
for k, v := range query {
|
||||
q.Add(k, v)
|
||||
}
|
||||
|
||||
// assign encoded query string to http request
|
||||
req.URL.RawQuery = q.Encode()
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
fmt.Println("Errored when sending request to the server")
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
_, err = ioutil.ReadAll(resp.Body) // drain the response
|
||||
return err
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue