From c294bc70dc5f2622c622bac60ef515bd956c4982 Mon Sep 17 00:00:00 2001 From: tgruben Date: Thu, 23 Feb 2023 18:35:52 -0600 Subject: [PATCH] limit memory for backuptar/restoretar (#2270) --- buffer/filebuffer.go | 110 +++++++++++++++++++++++++++++++ cmd/backup_tar.go | 1 + cmd/restore_tar.go | 2 + ctl/backup_tar.go | 150 ++++++++++++++----------------------------- ctl/restore_tar.go | 89 +++++++++++++++++-------- 5 files changed, 222 insertions(+), 130 deletions(-) create mode 100644 buffer/filebuffer.go diff --git a/buffer/filebuffer.go b/buffer/filebuffer.go new file mode 100644 index 000000000..3f0647465 --- /dev/null +++ b/buffer/filebuffer.go @@ -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 +} diff --git a/cmd/backup_tar.go b/cmd/backup_tar.go index 6c9612d77..641f4ee51 100644 --- a/cmd/backup_tar.go +++ b/cmd/backup_tar.go @@ -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 } diff --git a/cmd/restore_tar.go b/cmd/restore_tar.go index d626a488c..9a1cdc3fc 100644 --- a/cmd/restore_tar.go +++ b/cmd/restore_tar.go @@ -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, diff --git a/ctl/backup_tar.go b/ctl/backup_tar.go index 0b0780b0d..e8f65d90d 100644 --- a/ctl/backup_tar.go +++ b/ctl/backup_tar.go @@ -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 } diff --git a/ctl/restore_tar.go b/ctl/restore_tar.go index 80bf9c5e7..be1052f99 100644 --- a/ctl/restore_tar.go +++ b/ctl/restore_tar.go @@ -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 +}