Merge pull request #1119 from percona/PBM-1531-storage-opts-refactor

boris-ilijic · web-flow · commit 84813214223a · 2025-04-03T15:57:41.000+02:00
PBM-1531: Refactor storage api to use options for file size
diff --git a/cmd/pbm/restore.go b/cmd/pbm/restore.go
@@ -515,7 +515,7 @@ func runFinishRestore(o descrRestoreOpts, node string) (fmt.Stringer, error) {
 	path := fmt.Sprintf("%s/%s/cluster", defs.PhysRestoresDir, o.restore)
 	msg := outMsg{"Command sent. Check `pbm describe-restore ...` for the result."}
 	err = stg.Save(path+"."+string(defs.StatusCopyDone),
-		strings.NewReader(fmt.Sprintf("%d", time.Now().Unix())), -1)
+		strings.NewReader(fmt.Sprintf("%d", time.Now().Unix())))
 	return msg, err
 }
 
diff --git a/pbm/archive/conv.go b/pbm/archive/conv.go
@@ -51,7 +51,7 @@ func GenerateV1FromV2(ctx context.Context, stg storage.Storage, bcp, rs string)
 	}
 
 	metaV1Filename := fmt.Sprintf("%s/%s/%s", bcp, rs, MetaFile)
-	err = stg.Save(metaV1Filename, bytes.NewReader(data), int64(len(data)))
+	err = stg.Save(metaV1Filename, bytes.NewReader(data), storage.Size(int64(len(data))))
 	if err != nil {
 		return errors.Wrapf(err, "save %q", metaV1Filename)
 	}
diff --git a/pbm/backup/backup.go b/pbm/backup/backup.go
@@ -649,7 +649,7 @@ func writeMeta(stg storage.Storage, meta *BackupMeta) error {
 		return errors.Wrap(err, "marshal data")
 	}
 
-	err = stg.Save(meta.Name+defs.MetadataFileSuffix, bytes.NewReader(b), -1)
+	err = stg.Save(meta.Name+defs.MetadataFileSuffix, bytes.NewReader(b))
 	return errors.Wrap(err, "write to store")
 }
 
diff --git a/pbm/backup/logical.go b/pbm/backup/logical.go
@@ -170,7 +170,7 @@ func (b *Backup) doLogical(
 				return errors.Wrap(err, "get storage")
 			}
 			filepath := path.Join(bcp.Name, rsMeta.Name, ns+ext)
-			return stg.Save(filepath, r, nssSize[ns])
+			return stg.Save(filepath, r, storage.Size(nssSize[ns]))
 		},
 		bcp.Compression,
 		bcp.CompressionLevel)
diff --git a/pbm/restore/physical.go b/pbm/restore/physical.go
@@ -989,7 +989,7 @@ func (r *PhysRestore) Snapshot(
 		limit: 1 << 20, // 1Mb
 		write: func(name string, data io.Reader) error {
 			// Logger should be disabled due to: PBM-1531
-			return r.stg.Save(name, data, -1, storage.UseLogger(false))
+			return r.stg.Save(name, data, storage.UseLogger(false))
 		},
 	})
 	logger.PauseMgo()
diff --git a/pbm/storage/azure/azure.go b/pbm/storage/azure/azure.go
@@ -128,7 +128,7 @@ func (*Blob) Type() storage.Type {
 	return storage.Azure
 }
 
-func (b *Blob) Save(name string, data io.Reader, sizeb int64, options ...storage.Option) error {
+func (b *Blob) Save(name string, data io.Reader, options ...storage.Option) error {
 	opts := storage.GetDefaultOpts()
 	for _, opt := range options {
 		if err := opt(opts); err != nil {
@@ -137,8 +137,8 @@ func (b *Blob) Save(name string, data io.Reader, sizeb int64, options ...storage
 	}
 
 	bufsz := defaultUploadBuff
-	if sizeb > 0 {
-		ps := int(sizeb / maxBlocks * 11 / 10) // add 10% just in case
+	if opts.Size > 0 {
+		ps := int(opts.Size / maxBlocks * 11 / 10) // add 10% just in case
 		if ps > bufsz {
 			bufsz = ps
 		}
@@ -150,7 +150,7 @@ func (b *Blob) Save(name string, data io.Reader, sizeb int64, options ...storage
 	}
 
 	if b.log != nil && opts.UseLogger {
-		b.log.Debug("BufferSize is set to %d (~%dMb) | %d", bufsz, bufsz>>20, sizeb)
+		b.log.Debug("BufferSize is set to %d (~%dMb) | %d", bufsz, bufsz>>20, opts.Size)
 	}
 
 	_, err := b.c.UploadStream(context.TODO(),
diff --git a/pbm/storage/blackhole/blackhole.go b/pbm/storage/blackhole/blackhole.go
@@ -18,7 +18,7 @@ func (*Blackhole) Type() storage.Type {
 	return storage.Blackhole
 }
 
-func (*Blackhole) Save(_ string, data io.Reader, _ int64, _ ...storage.Option) error {
+func (*Blackhole) Save(_ string, data io.Reader, _ ...storage.Option) error {
 	_, err := io.Copy(io.Discard, data)
 	return err
 }
diff --git a/pbm/storage/fs/fs.go b/pbm/storage/fs/fs.go
@@ -151,7 +151,7 @@ func writeSync(finalpath string, data io.Reader) (err error) {
 	return nil
 }
 
-func (fs *FS) Save(name string, data io.Reader, _ int64, _ ...storage.Option) error {
+func (fs *FS) Save(name string, data io.Reader, _ ...storage.Option) error {
 	return writeSync(path.Join(fs.root, name), data)
 }
 
diff --git a/pbm/storage/gcs/gcs.go b/pbm/storage/gcs/gcs.go
@@ -142,7 +142,7 @@ func (*GCS) Type() storage.Type {
 	return storage.GCS
 }
 
-func (g *GCS) Save(name string, data io.Reader, size int64, _ ...storage.Option) error {
+func (g *GCS) Save(name string, data io.Reader, _ ...storage.Option) error {
 	ctx := context.Background()
 
 	w := g.bucketHandle.Object(path.Join(g.opts.Prefix, name)).NewWriter(ctx)
diff --git a/pbm/storage/s3/s3.go b/pbm/storage/s3/s3.go
@@ -314,7 +314,7 @@ func (*S3) Type() storage.Type {
 	return storage.S3
 }
 
-func (s *S3) Save(name string, data io.Reader, sizeb int64, options ...storage.Option) error {
+func (s *S3) Save(name string, data io.Reader, options ...storage.Option) error {
 	opts := storage.GetDefaultOpts()
 	for _, opt := range options {
 		if err := opt(opts); err != nil {
@@ -366,8 +366,8 @@ func (s *S3) Save(name string, data io.Reader, sizeb int64, options ...storage.O
 
 		partSize = int64(s.opts.UploadPartSize)
 	}
-	if sizeb > 0 {
-		ps := sizeb / int64(s.opts.MaxUploadParts) * 15 / 10 // add 50% just in case
+	if opts.Size > 0 {
+		ps := opts.Size / int64(s.opts.MaxUploadParts) * 15 / 10 // add 50% just in case
 		if ps > partSize {
 			partSize = ps
 		}
@@ -376,8 +376,8 @@ func (s *S3) Save(name string, data io.Reader, sizeb int64, options ...storage.O
 	if s.log != nil && opts.UseLogger {
 		s.log.Debug("uploading %q [size hint: %v (%v); part size: %v (%v)]",
 			name,
-			sizeb,
-			storage.PrettySize(sizeb),
+			opts.Size,
+			storage.PrettySize(opts.Size),
 			partSize,
 			storage.PrettySize(partSize))
 	}
diff --git a/pbm/storage/storage.go b/pbm/storage/storage.go
@@ -37,7 +37,7 @@ type FileInfo struct {
 
 type Storage interface {
 	Type() Type
-	Save(name string, data io.Reader, size int64, options ...Option) error
+	Save(name string, data io.Reader, options ...Option) error
 	SourceReader(name string) (io.ReadCloser, error)
 	// FileStat returns file info. It returns error if file is empty or not exists.
 	FileStat(name string) (FileInfo, error)
@@ -72,12 +72,14 @@ func ParseType(s string) Type {
 // Opts represents storage options
 type Opts struct {
 	UseLogger bool
+	Size      int64
 }
 
 // GetDefaultOpts creates default options.
 func GetDefaultOpts() *Opts {
 	return &Opts{
 		UseLogger: true,
+		Size:      -1,
 	}
 }
 
@@ -93,6 +95,18 @@ func UseLogger(useLogger bool) Option {
 	}
 }
 
+// Size option sets size of the file for the upload.
+// It's used to calculate params for multi-part upload.
+func Size(size int64) Option {
+	return func(o *Opts) error {
+		if size < 0 && size != -1 {
+			return errors.New("invalid size option value")
+		}
+		o.Size = size
+		return nil
+	}
+}
+
 // IsInitialized checks if there is PBM init file on the storage.
 func IsInitialized(ctx context.Context, stg Storage) (bool, error) {
 	_, err := stg.FileStat(defs.StorInitFile)
@@ -229,7 +243,7 @@ func Upload(
 
 	saveDone := make(chan struct{})
 	go func() {
-		rwErr.write = dst.Save(fname, r, sizeb)
+		rwErr.write = dst.Save(fname, r, Size(sizeb))
 		saveDone <- struct{}{}
 	}()
 
diff --git a/pbm/storage/storage_test_helpers.go b/pbm/storage/storage_test_helpers.go
@@ -19,7 +19,7 @@ func RunStorageTests(t *testing.T, stg Storage, stgType Type) {
 		name := "test.txt"
 		content := "content"
 
-		err := stg.Save(name, strings.NewReader(content), int64(len(content)))
+		err := stg.Save(name, strings.NewReader(content), Size(int64(len(content))))
 		if err != nil {
 			t.Fatalf("Save failed: %s", err)
 		}
@@ -45,7 +45,7 @@ func RunStorageTests(t *testing.T, stg Storage, stgType Type) {
 		}
 
 		for _, f := range filesToSave {
-			if err := stg.Save(f.name, strings.NewReader(f.content), int64(len(f.content))); err != nil {
+			if err := stg.Save(f.name, strings.NewReader(f.content), Size(int64(len(f.content)))); err != nil {
 				t.Fatalf("Save failed: %s", err)
 			}
 		}
@@ -75,7 +75,7 @@ func RunStorageTests(t *testing.T, stg Storage, stgType Type) {
 		name := "delete.txt"
 		content := "content"
 
-		if err := stg.Save(name, strings.NewReader(content), int64(len(content))); err != nil {
+		if err := stg.Save(name, strings.NewReader(content), Size(int64(len(content)))); err != nil {
 			t.Fatalf("Save failed: %s", err)
 		}
 
@@ -99,7 +99,7 @@ func RunStorageTests(t *testing.T, stg Storage, stgType Type) {
 		dst := "copydst.txt"
 		content := "copy content"
 
-		if err := stg.Save(src, strings.NewReader(content), int64(len(content))); err != nil {
+		if err := stg.Save(src, strings.NewReader(content), Size(int64(len(content)))); err != nil {
 			t.Fatalf("Save failed: %s", err)
 		}
 
@@ -121,7 +121,7 @@ func RunStorageTests(t *testing.T, stg Storage, stgType Type) {
 		name := "reader.txt"
 		content := "source reader content"
 
-		if err := stg.Save(name, strings.NewReader(content), int64(len(content))); err != nil {
+		if err := stg.Save(name, strings.NewReader(content), Size(int64(len(content)))); err != nil {
 			t.Fatalf("Save failed: %s", err)
 		}
 
diff --git a/pbm/util/storage.go b/pbm/util/storage.go
@@ -78,10 +78,10 @@ func Reinitialize(ctx context.Context, stg storage.Storage) error {
 }
 
 func RetryableWrite(stg storage.Storage, name string, data []byte) error {
-	err := stg.Save(name, bytes.NewBuffer(data), int64(len(data)))
+	err := stg.Save(name, bytes.NewBuffer(data), storage.Size(int64(len(data))))
 	if err != nil && stg.Type() == storage.Filesystem {
 		if fs.IsRetryableError(err) {
-			err = stg.Save(name, bytes.NewBuffer(data), int64(len(data)))
+			err = stg.Save(name, bytes.NewBuffer(data), storage.Size(int64(len(data))))
 		}
 	}
 

Original file line number	Diff line number	Diff line change
`@@ -515,7 +515,7 @@ func runFinishRestore(o descrRestoreOpts, node string) (fmt.Stringer, error) {`
`515`	`515`	`path := fmt.Sprintf("%s/%s/cluster", defs.PhysRestoresDir, o.restore)`
`516`	`516`	msg := outMsg{"Command sent. Check `pbm describe-restore ...` for the result."}
`517`	`517`	`err = stg.Save(path+"."+string(defs.StatusCopyDone),`
`518`		`- strings.NewReader(fmt.Sprintf("%d", time.Now().Unix())), -1)`
	`518`	`+ strings.NewReader(fmt.Sprintf("%d", time.Now().Unix())))`
`519`	`519`	`return msg, err`
`520`	`520`	`}`
`521`	`521`
Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ func GenerateV1FromV2(ctx context.Context, stg storage.Storage, bcp, rs string)`
`51`	`51`	`}`
`52`	`52`
`53`	`53`	`metaV1Filename := fmt.Sprintf("%s/%s/%s", bcp, rs, MetaFile)`
`54`		`- err = stg.Save(metaV1Filename, bytes.NewReader(data), int64(len(data)))`
	`54`	`+ err = stg.Save(metaV1Filename, bytes.NewReader(data), storage.Size(int64(len(data))))`
`55`	`55`	`if err != nil {`
`56`	`56`	`return errors.Wrapf(err, "save %q", metaV1Filename)`
`57`	`57`	`}`
Original file line number	Diff line number	Diff line change
`@@ -649,7 +649,7 @@ func writeMeta(stg storage.Storage, meta *BackupMeta) error {`
`649`	`649`	`return errors.Wrap(err, "marshal data")`
`650`	`650`	`}`
`651`	`651`
`652`		`- err = stg.Save(meta.Name+defs.MetadataFileSuffix, bytes.NewReader(b), -1)`
	`652`	`+ err = stg.Save(meta.Name+defs.MetadataFileSuffix, bytes.NewReader(b))`
`653`	`653`	`return errors.Wrap(err, "write to store")`
`654`	`654`	`}`
`655`	`655`
Original file line number	Diff line number	Diff line change
`@@ -170,7 +170,7 @@ func (b *Backup) doLogical(`
`170`	`170`	`return errors.Wrap(err, "get storage")`
`171`	`171`	`}`
`172`	`172`	`filepath := path.Join(bcp.Name, rsMeta.Name, ns+ext)`
`173`		`- return stg.Save(filepath, r, nssSize[ns])`
	`173`	`+ return stg.Save(filepath, r, storage.Size(nssSize[ns]))`
`174`	`174`	`},`
`175`	`175`	`bcp.Compression,`
`176`	`176`	`bcp.CompressionLevel)`
Original file line number	Diff line number	Diff line change
`@@ -128,7 +128,7 @@ func (*Blob) Type() storage.Type {`
`128`	`128`	`return storage.Azure`
`129`	`129`	`}`
`130`	`130`
`131`		`-func (b *Blob) Save(name string, data io.Reader, sizeb int64, options ...storage.Option) error {`
	`131`	`+func (b *Blob) Save(name string, data io.Reader, options ...storage.Option) error {`
`132`	`132`	`opts := storage.GetDefaultOpts()`
`133`	`133`	`for _, opt := range options {`
`134`	`134`	`if err := opt(opts); err != nil {`
`@@ -137,8 +137,8 @@ func (b *Blob) Save(name string, data io.Reader, sizeb int64, options ...storage`
`137`	`137`	`}`
`138`	`138`
`139`	`139`	`bufsz := defaultUploadBuff`
`140`		`- if sizeb > 0 {`
`141`		`- ps := int(sizeb / maxBlocks * 11 / 10) // add 10% just in case`
	`140`	`+ if opts.Size > 0 {`
	`141`	`+ ps := int(opts.Size / maxBlocks * 11 / 10) // add 10% just in case`
`142`	`142`	`if ps > bufsz {`
`143`	`143`	`bufsz = ps`
`144`	`144`	`}`
`@@ -150,7 +150,7 @@ func (b *Blob) Save(name string, data io.Reader, sizeb int64, options ...storage`
`150`	`150`	`}`
`151`	`151`
`152`	`152`	`if b.log != nil && opts.UseLogger {`
`153`		`- b.log.Debug("BufferSize is set to %d (~%dMb) \| %d", bufsz, bufsz>>20, sizeb)`
	`153`	`+ b.log.Debug("BufferSize is set to %d (~%dMb) \| %d", bufsz, bufsz>>20, opts.Size)`
`154`	`154`	`}`
`155`	`155`
`156`	`156`	`_, err := b.c.UploadStream(context.TODO(),`
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ func (*Blackhole) Type() storage.Type {`
`18`	`18`	`return storage.Blackhole`
`19`	`19`	`}`
`20`	`20`
`21`		`-func (*Blackhole) Save(_ string, data io.Reader, _ int64, _ ...storage.Option) error {`
	`21`	`+func (*Blackhole) Save(_ string, data io.Reader, _ ...storage.Option) error {`
`22`	`22`	`_, err := io.Copy(io.Discard, data)`
`23`	`23`	`return err`
`24`	`24`	`}`
Original file line number	Diff line number	Diff line change
`@@ -151,7 +151,7 @@ func writeSync(finalpath string, data io.Reader) (err error) {`
`151`	`151`	`return nil`
`152`	`152`	`}`
`153`	`153`
`154`		`-func (fs *FS) Save(name string, data io.Reader, _ int64, _ ...storage.Option) error {`
	`154`	`+func (fs *FS) Save(name string, data io.Reader, _ ...storage.Option) error {`
`155`	`155`	`return writeSync(path.Join(fs.root, name), data)`
`156`	`156`	`}`
`157`	`157`
Original file line number	Diff line number	Diff line change
`@@ -142,7 +142,7 @@ func (*GCS) Type() storage.Type {`
`142`	`142`	`return storage.GCS`
`143`	`143`	`}`
`144`	`144`
`145`		`-func (g *GCS) Save(name string, data io.Reader, size int64, _ ...storage.Option) error {`
	`145`	`+func (g *GCS) Save(name string, data io.Reader, _ ...storage.Option) error {`
`146`	`146`	`ctx := context.Background()`
`147`	`147`
`148`	`148`	`w := g.bucketHandle.Object(path.Join(g.opts.Prefix, name)).NewWriter(ctx)`