// PrioritizeWorkUnits changes the priorities of some number of work
// units. The actual work units are in options["work_unit_keys"]. A
// higher priority results in the work units being scheduled sooner.
func (jobs *JobServer) PrioritizeWorkUnits(workSpecName string, options map[string]interface{}) (bool, string, error) {
var (
err error
query coordinate.WorkUnitQuery
workSpec coordinate.WorkSpec
)
pwuOptions := PrioritizeWorkUnitsOptions{
Priority: math.NaN(),
Adjustment: math.NaN(),
}
workSpec, err = jobs.Namespace.WorkSpec(workSpecName)
if err == nil {
err = decode(&pwuOptions, options)
}
if err == nil && pwuOptions.WorkUnitKeys == nil {
return false, "missing work_unit_keys", err
}
if err == nil {
query.Names = pwuOptions.WorkUnitKeys
if !math.IsNaN(pwuOptions.Priority) {
err = workSpec.SetWorkUnitPriorities(query, pwuOptions.Priority)
} else if !math.IsNaN(pwuOptions.Adjustment) {
err = workSpec.AdjustWorkUnitPriorities(query, pwuOptions.Adjustment)
}
}
return err == nil, "", err
}
// WorkUnit creates a work unit in the named work spec; if it fails,
// fail the test.
func (a *CacheAssertions) WorkUnit(workSpec coordinate.WorkSpec, name string) coordinate.WorkUnit {
workUnit, err := workSpec.AddWorkUnit(name, map[string]interface{}{}, coordinate.WorkUnitMeta{})
if !a.NoError(err, "error creating work unit") {
a.FailNow("cannot create work unit")
}
return workUnit
}
// TestSpecCreateDestroy performs basic work spec lifetime tests.
func TestSpecCreateDestroy(t *testing.T) {
var (
dict map[string]interface{}
spec coordinate.WorkSpec
name string
names []string
err error
)
sts := SimpleTestSetup{
NamespaceName: "TestSpecCreateDestroy",
}
sts.SetUp(t)
defer sts.TearDown(t)
name = "spec"
spec, err = sts.Namespace.WorkSpec(name)
assert.Equal(t, coordinate.ErrNoSuchWorkSpec{Name: name}, err)
names, err = sts.Namespace.WorkSpecNames()
if assert.NoError(t, err) {
assert.Len(t, names, 0)
}
dict = map[string]interface{}{
"name": name,
"min_gb": 0.1,
}
spec, err = sts.Namespace.SetWorkSpec(dict)
if assert.NoError(t, err) && assert.NotNil(t, spec) {
assert.Equal(t, name, spec.Name())
}
spec, err = sts.Namespace.WorkSpec(name)
if assert.NoError(t, err) && assert.NotNil(t, spec) {
assert.Equal(t, name, spec.Name())
}
DataMatches(t, spec, dict)
names, err = sts.Namespace.WorkSpecNames()
if assert.NoError(t, err) {
assert.Equal(t, []string{name}, names)
}
err = sts.Namespace.DestroyWorkSpec(name)
assert.NoError(t, err)
spec, err = sts.Namespace.WorkSpec(name)
assert.Equal(t, coordinate.ErrNoSuchWorkSpec{Name: name}, err)
names, err = sts.Namespace.WorkSpecNames()
if assert.NoError(t, err) {
assert.Len(t, names, 0)
}
err = sts.Namespace.DestroyWorkSpec(name)
assert.Equal(t, coordinate.ErrNoSuchWorkSpec{Name: name}, err)
}
func (api *restAPI) fillWorkUnitShort(namespace coordinate.Namespace, spec coordinate.WorkSpec, name string, short *restdata.WorkUnitShort) error {
short.Name = name
return buildURLs(api.Router,
"namespace", namespace.Name(),
"spec", spec.Name(),
"unit", name,
).URL(&short.URL, "workUnit").Error
}
// ControlWorkSpec makes changes to a work spec that are not directly
// reflected in the work spec definition. This allows work specs to
// be paused or to stop generating new continuous jobs.
// ControlWorkSpecOptions has a complete listing of what can be done.
func (jobs *JobServer) ControlWorkSpec(workSpecName string, options map[string]interface{}) (bool, string, error) {
var (
cwsOptions ControlWorkSpecOptions
decoder *mapstructure.Decoder
err error
metadata mapstructure.Metadata
workSpec coordinate.WorkSpec
wsMeta coordinate.WorkSpecMeta
)
workSpec, err = jobs.Namespace.WorkSpec(workSpecName)
if err == nil {
// We care a lot about "false" vs. not present for
// these things. Manually create the decoder.
config := mapstructure.DecoderConfig{
Result: &cwsOptions,
Metadata: &metadata,
}
decoder, err = mapstructure.NewDecoder(&config)
}
if err == nil {
err = decoder.Decode(options)
}
// Get the existing metadata, then change it based on what
// we got provided
if err == nil {
wsMeta, err = workSpec.Meta(false)
}
if err == nil {
for _, key := range metadata.Keys {
switch key {
case "Continuous":
wsMeta.Continuous = cwsOptions.Continuous
case "Status":
wsMeta.Paused = cwsOptions.Status == Paused
case "Weight":
wsMeta.Weight = cwsOptions.Weight
case "Interval":
wsMeta.Interval = time.Duration(cwsOptions.Interval) * time.Second
case "MaxRunning":
wsMeta.MaxRunning = cwsOptions.MaxRunning
}
}
}
if err == nil {
err = workSpec.SetMeta(wsMeta)
}
return err == nil, "", err
}
func getWorkSpecData(spec coordinate.WorkSpec) (map[string]interface{}, error) {
data, err := spec.Data()
if err != nil {
return nil, err
}
byteify := func(key string) {
value, present := data[key]
if !present {
return
}
string, ok := value.(string)
if !ok {
return
}
data[key] = []byte(string)
}
byteify("module")
byteify("run_function")
return data, nil
}
// TestChainingMixed uses a combination of strings and tuples in its
// "output" data.
func TestChainingMixed(t *testing.T) {
var (
one, two coordinate.WorkSpec
attempt coordinate.Attempt
units map[string]coordinate.WorkUnit
err error
)
sts := SimpleTestSetup{
NamespaceName: "TestChainingMixed",
WorkerName: "worker",
}
sts.SetUp(t)
defer sts.TearDown(t)
one, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "one",
"then": "two",
})
if !assert.NoError(t, err) {
return
}
two, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "two",
})
if !assert.NoError(t, err) {
return
}
_, err = one.AddWorkUnit("a", map[string]interface{}{}, coordinate.WorkUnitMeta{})
assert.NoError(t, err)
sts.WorkSpec = one
attempt = sts.RequestOneAttempt(t)
err = attempt.Finish(map[string]interface{}{
"output": []interface{}{
"key",
cborrpc.PythonTuple{Items: []interface{}{
"key",
map[string]interface{}{
"data": "x",
},
map[string]interface{}{
"priority": 10.0,
},
}},
},
})
assert.NoError(t, err)
units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
if assert.NoError(t, err) {
if assert.Contains(t, units, "key") {
DataMatches(t, units["key"], map[string]interface{}{"data": "x"})
UnitHasPriority(t, units["key"], 10.0)
}
}
}
// TestChainingExpiry tests that, if an attempt finishes but is no
// longer the active attempt, then its successor work units will not
// be created.
func TestChainingExpiry(t *testing.T) {
var (
one, two coordinate.WorkSpec
err error
unit coordinate.WorkUnit
)
sts := SimpleTestSetup{
NamespaceName: "TestChainingExpiry",
WorkerName: "worker",
}
sts.SetUp(t)
defer sts.TearDown(t)
one, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "one",
"then": "two",
})
if !assert.NoError(t, err) {
return
}
sts.WorkSpec = one
two, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "two",
"disabled": true,
})
if !assert.NoError(t, err) {
return
}
// Create and perform a work unit, with no output
unit, err = one.AddWorkUnit("a", map[string]interface{}{}, coordinate.WorkUnitMeta{})
if !assert.NoError(t, err) {
return
}
attempt := sts.RequestOneAttempt(t)
// But wait! We got preempted
err = unit.ClearActiveAttempt()
assert.NoError(t, err)
sts.RequestOneAttempt(t)
// Now, let the original attempt finish, trying to generate
// more outputs
err = attempt.Finish(map[string]interface{}{
"output": []string{"unit"},
})
assert.NoError(t, err)
// Since attempt is no longer active, this shouldn't generate
// new outputs
units, err := two.WorkUnits(coordinate.WorkUnitQuery{})
if assert.NoError(t, err) {
assert.Empty(t, units)
}
}
func (api *restAPI) fillWorkUnit(namespace coordinate.Namespace, spec coordinate.WorkSpec, unit coordinate.WorkUnit, repr *restdata.WorkUnit) error {
err := api.fillWorkUnitShort(namespace, spec, unit.Name(), &repr.WorkUnitShort)
if err == nil {
repr.Data, err = unit.Data()
}
if err == nil {
var meta coordinate.WorkUnitMeta
meta, err = unit.Meta()
repr.Meta = &meta
}
if err == nil {
repr.Status, err = unit.Status()
}
if err == nil {
err = buildURLs(api.Router,
"namespace", namespace.Name(),
"spec", spec.Name(),
"unit", unit.Name(),
).
URL(&repr.WorkSpecURL, "workSpec").
URL(&repr.AttemptsURL, "workUnitAttempts").
Error
}
if err == nil {
var attempt coordinate.Attempt
attempt, err = unit.ActiveAttempt()
if err == nil && attempt != nil {
// This is cheating, a little, but it's probably
// the easiest way to reuse this code
var short restdata.AttemptShort
err = api.fillAttemptShort(namespace, attempt, &short)
if err == nil {
repr.ActiveAttemptURL = short.URL
}
}
}
return err
}
// GetWorkSpecMeta returns a set of control options for a given work
// spec. The returned map has the full set of keys that
// ControlWorkSpec() will accept.
func (jobs *JobServer) GetWorkSpecMeta(workSpecName string) (result map[string]interface{}, _ string, err error) {
var (
workSpec coordinate.WorkSpec
meta coordinate.WorkSpecMeta
)
workSpec, err = jobs.Namespace.WorkSpec(workSpecName)
if err == nil {
meta, err = workSpec.Meta(false)
}
if err == nil {
result = make(map[string]interface{})
if meta.Paused {
result["status"] = Paused
} else {
result["status"] = Runnable
}
result["continuous"] = meta.Continuous
result["interval"] = meta.Interval.Seconds()
result["max_running"] = meta.MaxRunning
result["weight"] = meta.Weight
}
return
}
// TestMetaContinuous specifically checks that you cannot enable the
// "continuous" flag on non-continuous work specs.
func TestMetaContinuous(t *testing.T) {
var (
err error
namespace coordinate.Namespace
spec coordinate.WorkSpec
meta coordinate.WorkSpecMeta
)
namespace, err = Coordinate.Namespace("TestMetaContinuous")
if !assert.NoError(t, err) {
return
}
defer namespace.Destroy()
spec, err = namespace.SetWorkSpec(map[string]interface{}{
"name": "spec",
"min_gb": 1,
})
if !assert.NoError(t, err) {
return
}
meta, err = spec.Meta(false)
if assert.NoError(t, err) {
assert.False(t, meta.Continuous)
assert.False(t, meta.CanBeContinuous)
}
meta.Continuous = true
err = spec.SetMeta(meta)
assert.NoError(t, err)
meta, err = spec.Meta(false)
if assert.NoError(t, err) {
// Cannot set the "continuous" flag
assert.False(t, meta.Continuous)
assert.False(t, meta.CanBeContinuous)
}
}
// ------------------------------------------------------------------------
// Coordinate setup helpers:
func createWorkUnits(spec coordinate.WorkSpec, n int, t assert.TestingT) {
for i := 0; i < n; i++ {
_, err := spec.AddWorkUnit(fmt.Sprintf("u%v", i), map[string]interface{}{}, coordinate.WorkUnitMeta{})
assert.NoError(t, err)
}
}
// TestSetMeta tests the basic SetMeta() call and a couple of its
// documented oddities.
func TestSetMeta(t *testing.T) {
var (
err error
namespace coordinate.Namespace
spec coordinate.WorkSpec
meta coordinate.WorkSpecMeta
)
namespace, err = Coordinate.Namespace("TestSetMeta")
if !assert.NoError(t, err) {
return
}
defer namespace.Destroy()
spec, err = namespace.SetWorkSpec(map[string]interface{}{
"name": "spec",
"min_gb": 1,
"continuous": true,
})
if !assert.NoError(t, err) {
return
}
meta, err = spec.Meta(false)
if assert.NoError(t, err) {
assert.Equal(t, 0, meta.Priority)
assert.Equal(t, 20, meta.Weight)
assert.False(t, meta.Paused)
assert.True(t, meta.Continuous)
assert.True(t, meta.CanBeContinuous)
assert.Zero(t, meta.Interval)
assert.WithinDuration(t, time.Time{}, meta.NextContinuous, 1*time.Microsecond)
assert.Equal(t, 0, meta.MaxRunning)
assert.Equal(t, 0, meta.MaxAttemptsReturned)
assert.Equal(t, "", meta.NextWorkSpecName)
assert.Equal(t, 0, meta.AvailableCount)
assert.Equal(t, 0, meta.PendingCount)
assert.Equal(t, "", meta.Runtime)
}
err = spec.SetMeta(coordinate.WorkSpecMeta{
Priority: 10,
Weight: 100,
Paused: true,
Continuous: false,
CanBeContinuous: false,
Interval: time.Duration(60) * time.Second,
MaxRunning: 10,
MaxAttemptsReturned: 1,
NextWorkSpecName: "then",
AvailableCount: 100,
PendingCount: 50,
Runtime: "go",
})
assert.NoError(t, err)
meta, err = spec.Meta(false)
if assert.NoError(t, err) {
assert.Equal(t, 10, meta.Priority)
assert.Equal(t, 100, meta.Weight)
assert.True(t, meta.Paused)
assert.False(t, meta.Continuous)
// Cannot clear "can be continuous" flag
assert.True(t, meta.CanBeContinuous)
assert.Equal(t, 60*time.Second, meta.Interval)
assert.WithinDuration(t, time.Time{}, meta.NextContinuous, 1*time.Microsecond)
assert.Equal(t, 10, meta.MaxRunning)
assert.Equal(t, 1, meta.MaxAttemptsReturned)
// Cannot change following work spec
assert.Equal(t, "", meta.NextWorkSpecName)
// Cannot set the counts
assert.Equal(t, 0, meta.AvailableCount)
assert.Equal(t, 0, meta.PendingCount)
// Cannot change the language runtime
assert.Equal(t, "", meta.Runtime)
}
}
// TestTwoWorkSpecsBasic ensures that two work specs can be created
// and have independent lifetimes.
func TestTwoWorkSpecsBasic(t *testing.T) {
var (
err error
dict1, dict2 map[string]interface{}
name1, name2 string
names []string
spec coordinate.WorkSpec
)
namespace, err := Coordinate.Namespace("TestTwoWorkSpecsBasic")
if !assert.NoError(t, err) {
return
}
defer namespace.Destroy()
name1 = "spec1"
name2 = "spec2"
dict1 = map[string]interface{}{"name": name1, "min_gb": 1}
dict2 = map[string]interface{}{"name": name2, "min_gb": 2}
names, err = namespace.WorkSpecNames()
if assert.NoError(t, err) {
assert.Len(t, names, 0)
}
spec, err = namespace.SetWorkSpec(dict1)
assert.NoError(t, err)
names, err = namespace.WorkSpecNames()
if assert.NoError(t, err) {
assert.Equal(t, []string{name1}, names)
}
spec, err = namespace.SetWorkSpec(dict2)
assert.NoError(t, err)
names, err = namespace.WorkSpecNames()
if assert.NoError(t, err) {
assert.Len(t, names, 2)
assert.Contains(t, names, name1)
assert.Contains(t, names, name2)
}
spec, err = namespace.WorkSpec(name1)
if assert.NoError(t, err) && assert.NotNil(t, spec) {
assert.Equal(t, name1, spec.Name())
}
spec, err = namespace.WorkSpec(name2)
if assert.NoError(t, err) && assert.NotNil(t, spec) {
assert.Equal(t, name2, spec.Name())
}
err = namespace.DestroyWorkSpec(name1)
assert.NoError(t, err)
names, err = namespace.WorkSpecNames()
if assert.NoError(t, err) {
assert.Len(t, names, 1)
assert.Contains(t, names, name2)
}
spec, err = namespace.WorkSpec(name1)
assert.Equal(t, coordinate.ErrNoSuchWorkSpec{Name: name1}, err)
spec, err = namespace.WorkSpec(name2)
if assert.NoError(t, err) && assert.NotNil(t, spec) {
assert.Equal(t, name2, spec.Name())
}
err = namespace.DestroyWorkSpec(name2)
assert.NoError(t, err)
names, err = namespace.WorkSpecNames()
if assert.NoError(t, err) {
assert.Len(t, names, 0)
}
spec, err = namespace.WorkSpec(name1)
assert.Equal(t, coordinate.ErrNoSuchWorkSpec{Name: name1}, err)
spec, err = namespace.WorkSpec(name2)
assert.Equal(t, coordinate.ErrNoSuchWorkSpec{Name: name2}, err)
}
// TestWorkUnitChaining tests that completing work units in one work spec
// will cause work units to appear in another, if so configured.
func TestWorkUnitChaining(t *testing.T) {
var (
err error
one, two coordinate.WorkSpec
units map[string]coordinate.WorkUnit
attempt coordinate.Attempt
)
sts := SimpleTestSetup{
NamespaceName: "TestWorkUnitChaining",
WorkerName: "worker",
}
sts.SetUp(t)
defer sts.TearDown(t)
one, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "one",
"then": "two",
})
if !assert.NoError(t, err) {
return
}
// RequestAttempts always returns this
sts.WorkSpec = one
two, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "two",
"disabled": true,
})
if !assert.NoError(t, err) {
return
}
// Create and perform a work unit, with no output
_, err = one.AddWorkUnit("a", map[string]interface{}{}, coordinate.WorkUnitMeta{})
assert.NoError(t, err)
sts.WorkSpec = one
attempt = sts.RequestOneAttempt(t)
err = attempt.Finish(nil)
assert.NoError(t, err)
units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
if assert.NoError(t, err) {
assert.Empty(t, units)
}
// Create and perform a work unit, with a map output
_, err = one.AddWorkUnit("b", map[string]interface{}{}, coordinate.WorkUnitMeta{})
assert.NoError(t, err)
attempt = sts.RequestOneAttempt(t)
err = attempt.Finish(map[string]interface{}{
"output": map[string]interface{}{
"two_b": map[string]interface{}{"k": "v"},
},
})
assert.NoError(t, err)
units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
if assert.NoError(t, err) {
assert.Len(t, units, 1)
if assert.Contains(t, units, "two_b") {
DataMatches(t, units["two_b"], map[string]interface{}{"k": "v"})
}
}
// Create and perform a work unit, with a slice output
_, err = one.AddWorkUnit("c", map[string]interface{}{}, coordinate.WorkUnitMeta{})
assert.NoError(t, err)
attempt = sts.RequestOneAttempt(t)
err = attempt.Finish(map[string]interface{}{
"output": []string{"two_c", "two_cc"},
})
assert.NoError(t, err)
units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
if assert.NoError(t, err) {
assert.Len(t, units, 3)
assert.Contains(t, units, "two_b")
assert.Contains(t, units, "two_cc")
if assert.Contains(t, units, "two_c") {
DataEmpty(t, units["two_c"])
}
}
// Put the output in the original work unit data
_, err = one.AddWorkUnit("d", map[string]interface{}{
"output": []string{"two_d"},
}, coordinate.WorkUnitMeta{})
assert.NoError(t, err)
attempt = sts.RequestOneAttempt(t)
err = attempt.Finish(nil)
assert.NoError(t, err)
units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
if assert.NoError(t, err) {
assert.Len(t, units, 4)
assert.Contains(t, units, "two_b")
assert.Contains(t, units, "two_c")
assert.Contains(t, units, "two_cc")
assert.Contains(t, units, "two_d")
}
}
// UpdateWorkUnit causes some state change in a work unit. If the
// work unit is pending, this is the principal interface to complete
// or renew it; if it is already complete this can cause it to be
// retried.
func (jobs *JobServer) UpdateWorkUnit(
workSpecName string,
workUnitKey string,
options map[string]interface{},
) (bool, string, error) {
// Note that in several corner cases, the behavior of this as
// written disagrees with Python coordinated's:
//
// * If neither "lease_time" nor "status" is specified,
// Python coordinated immediately returns False without
// checking if workUnitKey is valid
//
// * Python coordinated allows arbitrary status changes,
// including AVAILABLE -> FINISHED
//
// * This openly ignores "worker_id", as distinct from Python
// coordinated, which logs an obscure warning and changes it,
// but only on a renew
var (
attempt coordinate.Attempt
changed bool
err error
status coordinate.AttemptStatus
uwuOptions UpdateWorkUnitOptions
workSpec coordinate.WorkSpec
workUnit coordinate.WorkUnit
)
err = decode(&uwuOptions, options)
if err == nil {
workSpec, err = jobs.Namespace.WorkSpec(workSpecName)
}
if err == nil {
workUnit, err = workSpec.WorkUnit(workUnitKey)
}
if err == nil {
if workUnit == nil {
return false, fmt.Sprintf("no such work unit key=%v", workUnitKey), nil
}
}
if err == nil {
attempt, err = workUnit.ActiveAttempt()
}
if err == nil && attempt != nil {
status, err = attempt.Status()
}
if err == nil && attempt != nil {
if status == coordinate.Expired || status == coordinate.Retryable {
// The Python Coordinate API sees both of these
// statuses as "available", and we want to fall
// into the next block.
attempt = nil
}
}
if err == nil && attempt == nil {
// Caller is trying to manipulate an AVAILABLE work
// unit. Either they are trying to change the work
// unit data in place, or they are trying to jump a
// work unit directly to a completed state. (The
// latter is possible during the Python work unit
// parent cleanup, if the timing is bad.)
if uwuOptions.Status == Available || uwuOptions.Status == 0 {
// The only thing we are doing is changing the
// work unit data.
if uwuOptions.Data != nil {
meta, err := workUnit.Meta()
if err == nil {
_, err = workSpec.AddWorkUnit(workUnit.Name(), uwuOptions.Data, meta)
}
if err == nil {
changed = true
}
}
return changed && err == nil, "", err
}
// Otherwise we are trying to transition to another
// state; so force-create an attempt.
worker, err := jobs.Namespace.Worker(uwuOptions.WorkerID)
if err == nil {
attempt, err = worker.MakeAttempt(workUnit, uwuOptions.LeaseDuration())
status = coordinate.Pending
}
}
if err == nil {
switch status {
case coordinate.Pending:
changed = true // or there's an error
switch uwuOptions.Status {
case 0, Pending:
err = attempt.Renew(uwuOptions.LeaseDuration(), uwuOptions.Data)
case Available:
err = attempt.Expire(uwuOptions.Data)
case Finished:
err = attempt.Finish(uwuOptions.Data)
case Failed:
err = attempt.Fail(uwuOptions.Data)
default:
err = errors.New("update_work_unit invalid status")
}
case coordinate.Expired:
err = errors.New("update_work_unit logic error, trying to refresh expired unit")
case coordinate.Finished:
switch uwuOptions.Status {
case 0, Finished:
changed = false // no-op
case Available:
err = workUnit.ClearActiveAttempt()
changed = true
case Failed:
changed = false // see below
default:
err = errors.New("update_work_unit cannot change finished unit")
}
case coordinate.Failed:
switch uwuOptions.Status {
case 0, Failed:
changed = false // no-op
case Available: // "retry"
err = workUnit.ClearActiveAttempt()
changed = true
case Finished:
// The Python worker, with two separate
// processes, has a race wherein there
// could be 15 seconds to go, the parent
// kills off the child, and the child
// finishes successfully, all at the same
// time. In that case the successful
// finish should win.
err = attempt.Finish(nil)
changed = true
default:
err = errors.New("update_work_unit cannot change failed unit")
}
case coordinate.Retryable:
err = errors.New("update_work_unit logic error, trying to refresh retryable unit")
default:
err = fmt.Errorf("update_work_unit invalid attempt status %+v", status)
}
}
return changed && err == nil, "", err
}
// TestChainingTwoStep separately renews an attempt to insert an output
// key, then finishes the work unit; it should still chain.
func TestChainingTwoStep(t *testing.T) {
var (
one, two coordinate.WorkSpec
attempt coordinate.Attempt
units map[string]coordinate.WorkUnit
unit coordinate.WorkUnit
err error
)
sts := SimpleTestSetup{
NamespaceName: "TestChainingTwoStep",
WorkerName: "worker",
}
sts.SetUp(t)
defer sts.TearDown(t)
one, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "one",
"then": "two",
})
if !assert.NoError(t, err) {
return
}
two, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "two",
})
if !assert.NoError(t, err) {
return
}
_, err = one.AddWorkUnit("a", map[string]interface{}{}, coordinate.WorkUnitMeta{})
assert.NoError(t, err)
sts.WorkSpec = one
attempt = sts.RequestOneAttempt(t)
err = attempt.Renew(900*time.Second,
map[string]interface{}{
"output": []interface{}{
[]byte{1, 2, 3, 4},
cborrpc.PythonTuple{Items: []interface{}{
[]byte{1, 2, 3, 4},
map[interface{}]interface{}{},
map[interface{}]interface{}{
"priority": 0,
},
}},
},
})
assert.NoError(t, err)
err = attempt.Finish(nil)
assert.NoError(t, err)
units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
if assert.NoError(t, err) {
if assert.Contains(t, units, "\x01\x02\x03\x04") {
unit = units["\x01\x02\x03\x04"]
DataEmpty(t, unit)
UnitHasPriority(t, unit, 0.0)
}
}
}
// TestByRuntime creates two work specs with different runtimes, and
// validates that requests that want a specific runtime get it.
func TestByRuntime(t *testing.T) {
// The specific thing we'll simulate here is one Python
// worker, using the jobserver interface, with an empty
// runtime string, plus one Go worker, using the native API,
// with a "go" runtime.
var (
err error
pSpec, gSpec coordinate.WorkSpec
pUnit, gUnit coordinate.WorkUnit
attempts []coordinate.Attempt
)
sts := SimpleTestSetup{
NamespaceName: "TestByRuntime",
WorkerName: "worker",
}
sts.SetUp(t)
defer sts.TearDown(t)
pSpec, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "p",
})
if !assert.NoError(t, err) {
return
}
pUnit, err = pSpec.AddWorkUnit("p", map[string]interface{}{}, coordinate.WorkUnitMeta{})
if !assert.NoError(t, err) {
return
}
gSpec, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "g",
"runtime": "go",
})
if !assert.NoError(t, err) {
return
}
gUnit, err = gSpec.AddWorkUnit("g", map[string]interface{}{}, coordinate.WorkUnitMeta{})
if !assert.NoError(t, err) {
return
}
// If we use default settings for RequestAttempts, we should
// get back both work units
Clock.Add(5 * time.Second)
attempts, err = sts.Worker.RequestAttempts(coordinate.AttemptRequest{})
if assert.NoError(t, err) && assert.Len(t, attempts, 1) {
err = attempts[0].Finish(map[string]interface{}{})
assert.NoError(t, err)
wasP := attempts[0].WorkUnit().Name() == "p"
// Get more attempts
Clock.Add(time.Duration(5) * time.Second)
attempts, err = sts.Worker.RequestAttempts(coordinate.AttemptRequest{})
if assert.NoError(t, err) && assert.Len(t, attempts, 1) {
err = attempts[0].Finish(map[string]interface{}{})
assert.NoError(t, err)
// Should have gotten the other work spec
if wasP {
assert.Equal(t, "g", attempts[0].WorkUnit().Name())
} else {
assert.Equal(t, "p", attempts[0].WorkUnit().Name())
}
}
// Now there shouldn't be anything more
Clock.Add(5 * time.Second)
sts.RequestNoAttempts(t)
}
// Reset the world
err = pUnit.ClearActiveAttempt()
assert.NoError(t, err)
err = gUnit.ClearActiveAttempt()
assert.NoError(t, err)
// What we expect to get from jobserver
Clock.Add(5 * time.Second)
attempts, err = sts.Worker.RequestAttempts(coordinate.AttemptRequest{
Runtimes: []string{""},
})
if assert.NoError(t, err) && assert.Len(t, attempts, 1) {
assert.Equal(t, "p", attempts[0].WorkUnit().Name())
err = attempts[0].Retry(map[string]interface{}{}, time.Duration(0))
assert.NoError(t, err)
}
// A more sophisticated Python check
Clock.Add(5 * time.Second)
attempts, err = sts.Worker.RequestAttempts(coordinate.AttemptRequest{
Runtimes: []string{"python", "python_2", "python_2.7", ""},
})
if assert.NoError(t, err) && assert.Len(t, attempts, 1) {
assert.Equal(t, "p", attempts[0].WorkUnit().Name())
err = attempts[0].Retry(map[string]interface{}{}, time.Duration(0))
assert.NoError(t, err)
}
// What we expect to get from Go land
Clock.Add(5 * time.Second)
attempts, err = sts.Worker.RequestAttempts(coordinate.AttemptRequest{
Runtimes: []string{"go"},
})
if assert.NoError(t, err) && assert.Len(t, attempts, 1) {
assert.Equal(t, "g", attempts[0].WorkUnit().Name())
err = attempts[0].Retry(map[string]interface{}{}, time.Duration(0))
assert.NoError(t, err)
}
}
// TestChainingDuplicate tests that work unit chaining still works
// even when the same output work unit is generated twice (it should
// get retried).
func TestChainingDuplicate(t *testing.T) {
var (
err error
one, two coordinate.WorkSpec
attempt coordinate.Attempt
)
sts := SimpleTestSetup{
NamespaceName: "TestChainingDuplicate",
WorkerName: "worker",
}
sts.SetUp(t)
defer sts.TearDown(t)
one, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "one",
"then": "two",
"priority": 1,
})
if !assert.NoError(t, err) {
return
}
two, err = sts.Namespace.SetWorkSpec(map[string]interface{}{
"name": "two",
"priority": 2,
})
if !assert.NoError(t, err) {
return
}
_, err = one.AddWorkUnit("a", map[string]interface{}{}, coordinate.WorkUnitMeta{})
assert.NoError(t, err)
_, err = one.AddWorkUnit("b", map[string]interface{}{}, coordinate.WorkUnitMeta{})
assert.NoError(t, err)
sts.WorkSpec = one
attempt = sts.RequestOneAttempt(t)
assert.Equal(t, "a", attempt.WorkUnit().Name())
err = attempt.Finish(map[string]interface{}{
"output": []string{"z"},
})
assert.NoError(t, err)
sts.WorkSpec = two
attempt = sts.RequestOneAttempt(t)
assert.Equal(t, "z", attempt.WorkUnit().Name())
err = attempt.Finish(map[string]interface{}{})
assert.NoError(t, err)
sts.WorkSpec = one
attempt = sts.RequestOneAttempt(t)
assert.Equal(t, "b", attempt.WorkUnit().Name())
err = attempt.Finish(map[string]interface{}{
"output": []string{"z"},
})
assert.NoError(t, err)
sts.WorkSpec = two
attempt = sts.RequestOneAttempt(t)
assert.Equal(t, "z", attempt.WorkUnit().Name())
err = attempt.Finish(map[string]interface{}{})
assert.NoError(t, err)
sts.RequestNoAttempts(t)
}
