-
Notifications
You must be signed in to change notification settings - Fork 83
/
main.go
135 lines (131 loc) · 3.15 KB
/
main.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
131
132
133
134
135
package main
import (
"syscall"
// "io/ioutil"
"time"
"bytes"
"compress/zlib"
"crypto/md5"
"io"
"flag"
"sync"
"github.com/dropbox/rust-brotli/c/go/brotli"
)
var useWriter = flag.Bool("use_writer", false, "should use writer for brotli")
var parallelism = flag.Int("workers", 16, "number of workers to run with")
var threads = flag.Int("threads", 2, "number of threads to compress the file")
var iters = flag.Int("iter", 1024 * 1024, "number of times to run the compression/decompression cycle")
var sizeHint = flag.Uint64("size", 4096 * 1024, "Size of the test data")
var quality = flag.Float64("quality", 5, "brotli quality level")
var useZlib = flag.Int("zlib", 0, "1 for zlib 2 for both zlib and brotli")
var nongraceful = flag.Bool("hardexit", false, "use syscall.ExitGroup to end the program")
var anongoroutines = flag.Bool("anon", false, "use a separate anonymous goroutine for each invocation")
var timeout = flag.Duration("timeout", 0, "timeout until process exits with code 1")
func main() {
flag.Parse()
if *timeout != 0 {
go func() {
time.Sleep(*timeout)
syscall.Exit(1)
}()
}
options := brotli.CompressionOptions{
NumThreads: *threads,
Quality: float32(*quality),
Catable: true,
Appendable: true,
Magic: true,
SizeHint: uint32(*sizeHint),
}
file := testData(int(options.SizeHint))
if *anongoroutines {
var wg sync.WaitGroup
wg.Add(*parallelism)
for par :=0; par < *parallelism; par += 1 {
go func() {
for iter := 0; iter < *iters; iter += 1 {
_ = stresst(options, file, *useZlib % 2 == 0)
if *useZlib == 2 {
_ = stresst(options, file, false)
}
}
if *nongraceful {
syscall.Exit(0)
}
wg.Done()
}()
}
wg.Wait()
} else {
for iter := 0; iter < *iters; iter += 1 {
var wg sync.WaitGroup
wg.Add(*parallelism)
for par :=0; par < *parallelism; par += 1 {
go func() {
_ = stresst(options, file, *useZlib % 2 == 0)
if *useZlib == 2 {
_ = stresst(options, file, false)
}
wg.Done()
}()
}
wg.Wait()
}
if *nongraceful {
syscall.Exit(0)
}
}
}
func stresst(
options brotli.CompressionOptions,
file []byte,
useBrotli bool,
) (
md5out [md5.Size]byte,
) {
var output bytes.Buffer
input := bytes.NewBuffer(file)
var compressionWriter io.WriteCloser
var err error
if useBrotli {
compressionWriter = brotli.NewMultiCompressionWriter(
&output, options)
} else {
compressionWriter = zlib.NewWriter(&output)
}
if err != nil {
panic(err)
}
_, err = io.Copy(compressionWriter, input)
if err != nil {
panic(err)
}
err = compressionWriter.Close()
if err != nil {
panic(err)
}
md5out = md5.Sum(output.Bytes())
compressed := bytes.NewBuffer(output.Bytes())
var compressionReader io.ReadCloser
if useBrotli {
compressionReader = brotli.NewDecompressionReader(compressed)
} else {
compressionReader, err = zlib.NewReader(compressed)
}
if err != nil {
panic(err)
}
var rt bytes.Buffer
_, err = io.Copy(&rt, compressionReader)
if err != nil {
panic(err)
}
err = compressionReader.Close()
if err != nil {
panic(err)
}
if !bytes.Equal(rt.Bytes(), file) {
panic ("Files differ " + string(rt.Bytes()))
}
return
}