wal-g

1
package internal
2

3
import (
4
	"context"
5
	"io"
6

7
	"github.com/wal-g/wal-g/internal/ioextensions"
8
	"github.com/wal-g/wal-g/internal/limiters"
9
	"github.com/wal-g/wal-g/pkg/storages/storage"
10
	"golang.org/x/time/rate"
11
)
12

13
type LimitedFolder struct {
14
	storage.Folder
15
	limiter *rate.Limiter
16
}
17

18
func NewLimitedFolder(folder storage.Folder, limiter *rate.Limiter) *LimitedFolder {
19
	return &LimitedFolder{Folder: folder, limiter: limiter}
20
}
21

22
func (lf *LimitedFolder) GetSubFolder(subFolderRelativePath string) storage.Folder {
23
	folder := lf.Folder.GetSubFolder(subFolderRelativePath)
24
	return NewLimitedFolder(folder, lf.limiter)
25
}
26

27
func (lf *LimitedFolder) ReadObject(objectRelativePath string) (io.ReadCloser, error) {
28
	readCloser, err := lf.Folder.ReadObject(objectRelativePath)
29
	if err != nil {
30
		return nil, err
31
	}
32
	return ioextensions.ReadCascadeCloser{
33
		Reader: limiters.NewReader(context.Background(), readCloser, lf.limiter),
34
		Closer: readCloser,
35
	}, nil
36
}
37

38
func (lf *LimitedFolder) PutObject(name string, content io.Reader) error {
39
	return lf.PutObjectWithContext(context.Background(), name, content)
40
}
41

42
func (lf *LimitedFolder) PutObjectWithContext(ctx context.Context, name string, content io.Reader) error {
43
	limitedReader := limiters.NewReader(ctx, content, lf.limiter)
44
	return lf.Folder.PutObjectWithContext(ctx, name, limitedReader)
45
}
46