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
136
137
138
139
140
141
142
143
144
145
146
|
//go:build integration && perftest
// +build integration,perftest
package main
import (
"context"
"flag"
"fmt"
"io"
"io/ioutil"
"os"
"os/signal"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/credentials"
"github.com/aws/aws-sdk-go/aws/request"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/s3"
)
var config Config
func init() {
config.SetupFlags("", flag.CommandLine)
}
func main() {
if err := flag.CommandLine.Parse(os.Args[1:]); err != nil {
flag.CommandLine.PrintDefaults()
exitErrorf(err, "failed to parse CLI commands")
}
if err := config.Validate(); err != nil {
flag.CommandLine.PrintDefaults()
exitErrorf(err, "invalid arguments")
}
client := NewClient(config.Client)
var creds *credentials.Credentials
if config.SDK.Anonymous {
creds = credentials.AnonymousCredentials
}
var endpoint *string
if v := config.Endpoint; len(v) != 0 {
endpoint = &v
}
sess, err := session.NewSession(&aws.Config{
HTTPClient: client,
Endpoint: endpoint,
Credentials: creds,
S3Disable100Continue: aws.Bool(!config.SDK.ExpectContinue),
})
if err != nil {
exitErrorf(err, "failed to load config")
}
// Create context cancel for Ctrl+C/Interrupt
ctx, cancelFn := context.WithCancel(context.Background())
defer cancelFn()
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, os.Interrupt)
go func() {
<-sigCh
cancelFn()
}()
// Use the request duration timeout if specified.
if config.RequestDuration != 0 {
var timeoutFn func()
ctx, timeoutFn = context.WithTimeout(ctx, config.RequestDuration)
defer timeoutFn()
}
logger := NewLogger(os.Stdout)
// Start making the requests.
svc := s3.New(sess)
var reqCount int64
errCount := 0
for {
trace := doRequest(ctx, reqCount, svc, config)
select {
case <-ctx.Done():
return
default:
}
logger.RecordTrace(trace)
if err := trace.Err(); err != nil {
fmt.Fprintf(os.Stderr, err.Error())
errCount++
} else {
errCount = 0
}
if config.RequestCount > 0 && reqCount == config.RequestCount {
return
}
reqCount++
// If the first several requests fail, exist, something is broken.
if errCount == 5 && reqCount == 5 {
exitErrorf(trace.Err(), "unable to make requests")
}
if config.RequestDelay > 0 {
time.Sleep(config.RequestDelay)
}
}
}
func doRequest(ctx context.Context, id int64, svc *s3.S3, config Config) *RequestTrace {
traceCtx := NewRequestTrace(ctx, id)
defer traceCtx.RequestDone()
resp, err := svc.GetObjectWithContext(traceCtx, &s3.GetObjectInput{
Bucket: &config.Bucket,
Key: &config.Key,
}, func(r *request.Request) {
r.Handlers.Send.PushFront(traceCtx.OnSendAttempt)
r.Handlers.Complete.PushBack(traceCtx.OnCompleteRequest)
r.Handlers.CompleteAttempt.PushBack(traceCtx.OnCompleteAttempt)
})
if err != nil {
traceCtx.AppendError(fmt.Errorf("request failed, %v", err))
return traceCtx
}
defer resp.Body.Close()
if n, err := io.Copy(ioutil.Discard, resp.Body); err != nil {
traceCtx.AppendError(fmt.Errorf("read request body failed, read %v, %v", n, err))
return traceCtx
}
return traceCtx
}
func exitErrorf(err error, msg string, args ...interface{}) {
fmt.Fprintf(os.Stderr, "FAILED: %v\n"+msg+"\n", append([]interface{}{err}, args...)...)
os.Exit(1)
}
|
