-
Notifications
You must be signed in to change notification settings - Fork 567
/
file.go
130 lines (113 loc) · 2.67 KB
/
file.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
package pfsload
import (
"fmt"
"io"
"math/rand"
"path"
"github.com/pachyderm/pachyderm/v2/src/internal/errors"
"github.com/pachyderm/pachyderm/v2/src/internal/randutil"
)
const pathSize = 32
type RandomFile struct {
path string
r io.Reader
}
func NewRandomFile(path string, r io.Reader) *RandomFile {
return &RandomFile{
path: path,
r: r,
}
}
func (f *RandomFile) Path() string {
return f.path
}
func (f *RandomFile) Read(data []byte) (int, error) {
res, err := f.r.Read(data)
return res, errors.EnsureStack(err)
}
type FileSource interface {
Next() (*RandomFile, error)
}
func NewFileSource(spec *FileSourceSpec, random *rand.Rand) FileSource {
return newRandomFileSource(spec.Random, random)
}
type randomFileSource struct {
spec *RandomFileSourceSpec
random *rand.Rand
dirSource *randomDirectorySource
next int64
}
func newRandomFileSource(spec *RandomFileSourceSpec, random *rand.Rand) FileSource {
var dirSource *randomDirectorySource
if spec.Directory != nil {
dirSource = &randomDirectorySource{
spec: spec.Directory,
random: random,
}
}
return &randomFileSource{
spec: spec,
random: random,
dirSource: dirSource,
}
}
func (rfs *randomFileSource) Next() (*RandomFile, error) {
sizeSpec, err := FuzzSize(rfs.spec.Sizes, rfs.random)
if err != nil {
return nil, err
}
min, max := sizeSpec.MinSize, sizeSpec.MaxSize
size := min
if max > min {
size += rfs.random.Int63n(max - min)
}
return NewRandomFile(rfs.nextPath(), randutil.NewBytesReader(rfs.random, int64(size))), nil
}
func (rfs *randomFileSource) nextPath() string {
var dir string
if rfs.dirSource != nil {
dir = rfs.dirSource.nextPath()
}
if rfs.spec.IncrementPath {
next := rfs.next
rfs.next += 1
return path.Join(dir, fmt.Sprintf("%016d", next))
}
return path.Join(dir, string(randutil.Bytes(rfs.random, pathSize)))
}
type randomDirectorySource struct {
spec *RandomDirectorySpec
random *rand.Rand
next string
run int64
}
func (rds *randomDirectorySource) nextPath() string {
if rds.next == "" {
min, max := rds.spec.Depth.MinSize, rds.spec.Depth.MaxSize
depth := int(min)
if max > min {
depth += rds.random.Intn(int(max - min))
}
for i := 0; i < depth; i++ {
rds.next = path.Join(rds.next, string(randutil.Bytes(rds.random, pathSize)))
}
}
dir := rds.next
rds.run++
if rds.run == rds.spec.Run {
rds.next = ""
rds.run = 0
}
return dir
}
func Files(fileSource FileSource, count int) ([]*RandomFile, error) {
var files []*RandomFile
for i := 0; i < count; i++ {
file, err := fileSource.Next()
if err != nil {
return nil, errors.EnsureStack(err)
}
files = append(files, file)
}
return files, nil
}