-
Notifications
You must be signed in to change notification settings - Fork 132
/
submitter.go
137 lines (127 loc) · 3.58 KB
/
submitter.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
136
137
package submitter
import (
"context"
"sync"
"time"
"github.com/pkg/errors"
"github.com/armadaproject/armada/internal/common/armadaerrors"
"github.com/armadaproject/armada/pkg/api"
"github.com/armadaproject/armada/pkg/client"
)
type Submitter struct {
ApiConnectionDetails *client.ApiConnectionDetails
// Jobs to submit.
Queue string
JobSetName string
Jobs []*api.JobSubmitRequestItem
// If true, generate random clientId for Job, if not already provided
RandomClientId bool
// Number of batches of jobs to submit.
// A value of 0 indicates infinity.
NumBatches uint32
// Number of copies of the provided job to submit per batch.
BatchSize uint32
// Time between batches.
Interval time.Duration
// Number of seconds to wait for jobs to finish.
Timeout time.Duration
jobIds []string
mu sync.Mutex
}
func NewSubmitterFromTestSpec(conn *client.ApiConnectionDetails, testSpec *api.TestSpec) *Submitter {
return &Submitter{
ApiConnectionDetails: conn,
Jobs: testSpec.Jobs,
Queue: testSpec.Queue,
JobSetName: testSpec.JobSetId,
NumBatches: testSpec.NumBatches,
BatchSize: testSpec.BatchSize,
Interval: testSpec.Interval,
RandomClientId: testSpec.RandomClientId,
}
}
func (config *Submitter) Validate() error {
if len(config.Jobs) == 0 {
return errors.WithStack(&armadaerrors.ErrInvalidArgument{
Name: "Jobs",
Value: config.Jobs,
Message: "no jobs provided",
})
}
if config.Queue == "" {
return errors.WithStack(&armadaerrors.ErrInvalidArgument{
Name: "Queue",
Value: config.Queue,
Message: "not provided",
})
}
if config.JobSetName == "" {
return errors.WithStack(&armadaerrors.ErrInvalidArgument{
Name: "JobSetName",
Value: config.JobSetName,
Message: "not provided",
})
}
if config.BatchSize <= 0 {
return errors.WithStack(&armadaerrors.ErrInvalidArgument{
Name: "BatchSize",
Value: config.BatchSize,
Message: "batch size must be positive",
})
}
return nil
}
func (srv *Submitter) Run(ctx context.Context) error {
var numBatchesSent uint32
req := &api.JobSubmitRequest{
Queue: srv.Queue,
JobSetId: srv.JobSetName,
}
for i := 0; i < int(srv.BatchSize); i++ {
// workaround to always create new structs instead of copying the pointer
// that way, we can change some job values without affecting other jobs
for _, job := range srv.Jobs {
cloned := *job
req.JobRequestItems = append(req.JobRequestItems, &cloned)
}
}
if srv.RandomClientId {
client.AddClientIds(req.JobRequestItems)
}
return client.WithSubmitClient(srv.ApiConnectionDetails, func(c api.SubmitClient) error {
// Create a closed ticker channel; receiving on tickerCh returns immediately.
C := make(chan time.Time)
close(C)
tickerCh := (<-chan time.Time)(C)
// If an interval is provided, replace tickerCh with one that generates ticks periodically.
if srv.Interval != 0 {
ticker := time.NewTicker(srv.Interval)
defer ticker.Stop()
tickerCh = ticker.C
}
// Submit jobs.
for srv.NumBatches == 0 || numBatchesSent < srv.NumBatches {
select {
case <-ctx.Done():
return ctx.Err()
case <-tickerCh:
res, err := c.SubmitJobs(ctx, req)
if err != nil {
return errors.WithMessage(err, "error submitting jobs")
}
srv.mu.Lock()
for _, item := range res.JobResponseItems {
srv.jobIds = append(srv.jobIds, item.JobId)
}
srv.mu.Unlock()
numBatchesSent++
}
}
return nil
})
}
func (srv *Submitter) JobIds() []string {
srv.mu.Lock()
defer srv.mu.Unlock()
return srv.jobIds[:]
}