Golang WorkSpec.AddWorkUnit方法代码示例

// TestChainingExpiry tests that, if an attempt finishes but is no
// longer the active attempt, then its successor work units will not
// be created.
func (s *Suite) TestChainingExpiry(c *check.C) {
	var (
		one, two coordinate.WorkSpec
		err      error
		worker   coordinate.Worker
		unit     coordinate.WorkUnit
		attempts []coordinate.Attempt
	)

	one, err = s.Namespace.SetWorkSpec(map[string]interface{}{
		"name": "one",
		"then": "two",
	})
	c.Assert(err, check.IsNil)

	two, err = s.Namespace.SetWorkSpec(map[string]interface{}{
		"name":     "two",
		"disabled": true,
	})
	c.Assert(err, check.IsNil)

	worker, err = s.Namespace.Worker("worker")
	c.Assert(err, check.IsNil)

	// Create and perform a work unit, with no output
	unit, err = one.AddWorkUnit("a", map[string]interface{}{}, 0.0)
	c.Assert(err, check.IsNil)

	attempts, err = worker.RequestAttempts(coordinate.AttemptRequest{})
	c.Assert(err, check.IsNil)
	c.Assert(attempts, check.HasLen, 1)
	attempt := attempts[0]

	// But wait!  We got preempted
	err = unit.ClearActiveAttempt()
	c.Assert(err, check.IsNil)
	attempts, err = worker.RequestAttempts(coordinate.AttemptRequest{})
	c.Assert(err, check.IsNil)
	c.Assert(attempts, check.HasLen, 1)

	// Now, let the original attempt finish, trying to generate
	// more outputs
	err = attempt.Finish(map[string]interface{}{
		"output": []string{"unit"},
	})
	c.Assert(err, check.IsNil)

	// Since attempt is no longer active, this shouldn't generate
	// new outputs
	units, err := two.WorkUnits(coordinate.WorkUnitQuery{})
	c.Assert(err, check.IsNil)
	c.Check(units, check.HasLen, 0)
}

// TestTrivialWorkUnitFlow tests work unit creation, deletion, and existence.
func (s *Suite) TestTrivialWorkUnitFlow(c *check.C) {
	var (
		count int
		err   error
		spec  coordinate.WorkSpec
		unit  coordinate.WorkUnit
		units map[string]coordinate.WorkUnit
	)

	spec, err = s.Namespace.SetWorkSpec(map[string]interface{}{
		"name":   "spec",
		"min_gb": 1,
	})
	c.Assert(err, check.IsNil)

	unit, err = spec.AddWorkUnit("unit", map[string]interface{}{}, 0)
	c.Assert(err, check.IsNil)
	c.Check(unit.Name(), check.Equals, "unit")
	c.Check(unit.WorkSpec().Name(), check.Equals, "spec")

	unit, err = spec.WorkUnit("unit")
	c.Assert(err, check.IsNil)
	c.Check(unit.Name(), check.Equals, "unit")
	c.Check(unit.WorkSpec().Name(), check.Equals, "spec")

	units, err = spec.WorkUnits(coordinate.WorkUnitQuery{})
	c.Assert(err, check.IsNil)
	c.Check(units, check.HasLen, 1)
	c.Check(units["unit"], check.NotNil)
	c.Check(units["unit"].Name(), check.Equals, "unit")
	c.Check(units["unit"].WorkSpec().Name(), check.Equals, "spec")

	count, err = spec.DeleteWorkUnits(coordinate.WorkUnitQuery{})
	c.Assert(err, check.IsNil)
	c.Check(count, check.Equals, 1)

	unit, err = spec.WorkUnit("unit")
	c.Assert(err, check.IsNil)
	c.Check(unit, check.IsNil)

	units, err = spec.WorkUnits(coordinate.WorkUnitQuery{})
	c.Assert(err, check.IsNil)
	c.Check(units, check.HasLen, 0)
}

// makeWorkUnits creates a handful of work units within a work spec.
// These have keys "available", "pending", "finished", "failed",
// "expired", and "retryable", and wind up in the corresponding
// states.
func makeWorkUnits(spec coordinate.WorkSpec, worker coordinate.Worker) (map[string]coordinate.WorkUnit, error) {
	result := map[string]coordinate.WorkUnit{
		"available": nil,
		"pending":   nil,
		"finished":  nil,
		"failed":    nil,
		"expired":   nil,
		"retryable": nil,
	}
	for key := range result {
		unit, err := spec.AddWorkUnit(key, map[string]interface{}{}, 0)
		if err != nil {
			return nil, err
		}
		result[key] = unit

		// Run the workflow
		if key == "available" {
			continue
		}
		attempt, err := worker.MakeAttempt(unit, time.Duration(0))
		if err != nil {
			return nil, err
		}
		switch key {
		case "pending":
			{
			} // leave it running
		case "finished":
			err = attempt.Finish(nil)
		case "failed":
			err = attempt.Fail(nil)
		case "expired":
			err = attempt.Expire(nil)
		case "retryable":
			err = attempt.Retry(nil)
		}
		if err != nil {
			return nil, err
		}
	}
	return result, nil
}

// ------------------------------------------------------------------------
// Coordinate setup helpers:
func createWorkUnits(spec coordinate.WorkSpec, n int, c *check.C) {
	for i := 0; i < n; i++ {
		_, err := spec.AddWorkUnit(fmt.Sprintf("u%v", i), map[string]interface{}{}, 0.0)
		c.Assert(err, check.IsNil)
	}
}

// TestChainingTwoStep separately renews an attempt to insert an output
// key, then finishes the work unit; it should still chain.
func (s *Suite) TestChainingTwoStep(c *check.C) {
	var (
		one, two coordinate.WorkSpec
		worker   coordinate.Worker
		attempts []coordinate.Attempt
		units    map[string]coordinate.WorkUnit
		unit     coordinate.WorkUnit
		data     map[string]interface{}
		priority float64
		ok       bool
		err      error
	)

	one, err = s.Namespace.SetWorkSpec(map[string]interface{}{
		"name": "one",
		"then": "two",
	})
	c.Assert(err, check.IsNil)

	two, err = s.Namespace.SetWorkSpec(map[string]interface{}{
		"name": "two",
	})
	c.Assert(err, check.IsNil)

	worker, err = s.Namespace.Worker("worker")
	c.Assert(err, check.IsNil)

	_, err = one.AddWorkUnit("a", map[string]interface{}{}, 0.0)
	c.Assert(err, check.IsNil)

	attempts, err = worker.RequestAttempts(coordinate.AttemptRequest{})
	c.Assert(err, check.IsNil)
	c.Assert(attempts, check.HasLen, 1)

	err = attempts[0].Renew(time.Duration(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,
					},
				}},
			},
		})
	c.Assert(err, check.IsNil)

	err = attempts[0].Finish(nil)

	units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
	c.Assert(err, check.IsNil)
	c.Check(units, HasKeys, []string{"\x01\x02\x03\x04"})
	if unit, ok = units["\x01\x02\x03\x04"]; ok {
		data, err = unit.Data()
		c.Assert(err, check.IsNil)
		c.Check(data, check.DeepEquals, map[string]interface{}{})

		priority, err = unit.Priority()
		c.Assert(err, check.IsNil)
		c.Check(priority, check.Equals, 0.0)
	}
}

// TestChainingMixed uses a combination of strings and tuples in its
// "output" data.
func (s *Suite) TestChainingMixed(c *check.C) {
	var (
		one, two coordinate.WorkSpec
		worker   coordinate.Worker
		attempts []coordinate.Attempt
		units    map[string]coordinate.WorkUnit
		unit     coordinate.WorkUnit
		data     map[string]interface{}
		priority float64
		ok       bool
		err      error
	)

	one, err = s.Namespace.SetWorkSpec(map[string]interface{}{
		"name": "one",
		"then": "two",
	})
	c.Assert(err, check.IsNil)

	two, err = s.Namespace.SetWorkSpec(map[string]interface{}{
		"name": "two",
	})
	c.Assert(err, check.IsNil)

	worker, err = s.Namespace.Worker("worker")
	c.Assert(err, check.IsNil)

	_, err = one.AddWorkUnit("a", map[string]interface{}{}, 0.0)
	c.Assert(err, check.IsNil)

	attempts, err = worker.RequestAttempts(coordinate.AttemptRequest{})
	c.Assert(err, check.IsNil)
	c.Assert(attempts, check.HasLen, 1)

	err = attempts[0].Finish(map[string]interface{}{
		"output": []interface{}{
			"key",
			cborrpc.PythonTuple{Items: []interface{}{
				"key",
				map[string]interface{}{
					"data": "x",
				},
				map[string]interface{}{
					"priority": 10.0,
				},
			}},
		},
	})
	c.Assert(err, check.IsNil)

	units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
	c.Assert(err, check.IsNil)
	c.Check(units, HasKeys, []string{"key"})
	if unit, ok = units["key"]; ok {
		data, err = unit.Data()
		c.Assert(err, check.IsNil)
		c.Check(data, check.DeepEquals, map[string]interface{}{"data": "x"})

		priority, err = unit.Priority()
		c.Assert(err, check.IsNil)
		c.Check(priority, check.Equals, 10.0)
	}
}

// TestWorkUnitChaining tests that completing work units in one work spec
// will cause work units to appear in another, if so configured.
func (s *Suite) TestWorkUnitChaining(c *check.C) {
	var (
		err      error
		worker   coordinate.Worker
		one, two coordinate.WorkSpec
		units    map[string]coordinate.WorkUnit
		attempts []coordinate.Attempt
		data     map[string]interface{}
		ok       bool
	)

	one, err = s.Namespace.SetWorkSpec(map[string]interface{}{
		"name": "one",
		"then": "two",
	})
	c.Assert(err, check.IsNil)

	two, err = s.Namespace.SetWorkSpec(map[string]interface{}{
		"name":     "two",
		"disabled": true,
	})
	c.Assert(err, check.IsNil)

	worker, err = s.Namespace.Worker("worker")
	c.Assert(err, check.IsNil)

	// Create and perform a work unit, with no output
	_, err = one.AddWorkUnit("a", map[string]interface{}{}, 0.0)
	c.Assert(err, check.IsNil)

	attempts, err = worker.RequestAttempts(coordinate.AttemptRequest{})
	c.Assert(err, check.IsNil)
	c.Assert(attempts, check.HasLen, 1)

	err = attempts[0].Finish(nil)
	c.Assert(err, check.IsNil)

	units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
	c.Assert(err, check.IsNil)
	c.Check(units, HasKeys, []string{})

	// Create and perform a work unit, with a map output
	_, err = one.AddWorkUnit("b", map[string]interface{}{}, 0.0)
	c.Assert(err, check.IsNil)

	attempts, err = worker.RequestAttempts(coordinate.AttemptRequest{})
	c.Assert(err, check.IsNil)
	c.Assert(attempts, check.HasLen, 1)

	err = attempts[0].Finish(map[string]interface{}{
		"output": map[string]interface{}{
			"two_b": map[string]interface{}{"k": "v"},
		},
	})
	c.Assert(err, check.IsNil)

	units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
	c.Assert(err, check.IsNil)
	c.Check(units, HasKeys, []string{"two_b"})

	if _, ok = units["two_b"]; ok {
		data, err = units["two_b"].Data()
		c.Assert(err, check.IsNil)
		c.Check(data, check.DeepEquals, map[string]interface{}{"k": "v"})
	}

	// Create and perform a work unit, with a slice output
	_, err = one.AddWorkUnit("c", map[string]interface{}{}, 0.0)
	c.Assert(err, check.IsNil)

	attempts, err = worker.RequestAttempts(coordinate.AttemptRequest{})
	c.Assert(err, check.IsNil)
	c.Assert(attempts, check.HasLen, 1)

	err = attempts[0].Finish(map[string]interface{}{
		"output": []string{"two_c", "two_cc"},
	})
	c.Assert(err, check.IsNil)

	units, err = two.WorkUnits(coordinate.WorkUnitQuery{})
	c.Assert(err, check.IsNil)
	c.Check(units, HasKeys, []string{"two_b", "two_c", "two_cc"})

	if _, ok = units["two_c"]; ok {
		data, err = units["two_c"].Data()
		c.Assert(err, check.IsNil)
		c.Check(data, check.DeepEquals, map[string]interface{}{})
	}

	// Put the output in the original work unit data
	_, err = one.AddWorkUnit("d", map[string]interface{}{
		"output": []string{"two_d"},
	}, 0.0)
	c.Assert(err, check.IsNil)
	attempts, err = worker.RequestAttempts(coordinate.AttemptRequest{})
	c.Assert(err, check.IsNil)
	c.Assert(attempts, check.HasLen, 1)

//.........这里部分代码省略.........

// TestAttemptLifetime validates a basic attempt lifetime.
func (s *Suite) TestAttemptLifetime(c *check.C) {
	var (
		err               error
		data              map[string]interface{}
		attempt, attempt2 coordinate.Attempt
		aStatus           coordinate.AttemptStatus
		spec              coordinate.WorkSpec
		unit              coordinate.WorkUnit
		worker            coordinate.Worker
		uStatus           coordinate.WorkUnitStatus
	)
	spec, worker = s.makeWorkSpecAndWorker(c)

	// Create a work unit
	unit, err = spec.AddWorkUnit("a", map[string]interface{}{}, 0.0)
	c.Assert(err, check.IsNil)

	// The work unit should be "available"
	uStatus, err = unit.Status()
	c.Assert(err, check.IsNil)
	c.Check(uStatus, check.Equals, coordinate.AvailableUnit)

	// The work unit data should be defined but empty
	data, err = unit.Data()
	c.Assert(err, check.IsNil)
	c.Check(data, check.HasLen, 0)

	// Get an attempt for it
	attempts, err := worker.RequestAttempts(coordinate.AttemptRequest{})
	c.Assert(err, check.IsNil)
	c.Assert(attempts, check.HasLen, 1)
	attempt = attempts[0]

	// The work unit should be "pending"
	uStatus, err = unit.Status()
	c.Assert(err, check.IsNil)
	c.Check(uStatus, check.Equals, coordinate.PendingUnit)

	// The attempt should be "pending" too
	aStatus, err = attempt.Status()
	c.Assert(err, check.IsNil)
	c.Check(aStatus, check.Equals, coordinate.Pending)

	// The active attempt for the unit should match this
	attempt2, err = unit.ActiveAttempt()
	c.Assert(err, check.IsNil)
	c.Check(attempt2, AttemptMatches, attempt)

	// There should be one active attempt for the worker and it should
	// also match
	attempts, err = worker.ActiveAttempts()
	c.Assert(err, check.IsNil)
	c.Check(attempts, check.HasLen, 1)
	if len(attempts) > 0 {
		c.Check(attempts[0], AttemptMatches, attempt)
	}

	// The work unit data should (still) be defined but empty
	data, err = unit.Data()
	c.Assert(err, check.IsNil)
	c.Check(data, check.HasLen, 0)

	// Now finish the attempt with some updated data
	err = attempt.Finish(map[string]interface{}{
		"outputs": []string{"yes"},
	})
	c.Assert(err, check.IsNil)

	// The unit should report "finished"
	uStatus, err = unit.Status()
	c.Assert(err, check.IsNil)
	c.Check(uStatus, check.Equals, coordinate.FinishedUnit)

	// The attempt should report "finished"
	aStatus, err = attempt.Status()
	c.Assert(err, check.IsNil)
	c.Check(aStatus, check.Equals, coordinate.Finished)

	// The attempt should still be the active attempt for the unit
	attempt2, err = unit.ActiveAttempt()
	c.Assert(err, check.IsNil)
	c.Check(attempt2, AttemptMatches, attempt)

	// The attempt should not be in the active attempt list for the worker
	attempts, err = worker.ActiveAttempts()
	c.Assert(err, check.IsNil)
	c.Check(attempts, check.HasLen, 0)

	// Both the unit and the worker should have one archived attempt
	attempts, err = unit.Attempts()
	c.Assert(err, check.IsNil)
	c.Check(attempts, check.HasLen, 1)
	if len(attempts) > 0 {
		c.Check(attempts[0], AttemptMatches, attempt)
	}

	attempts, err = worker.AllAttempts()
	c.Assert(err, check.IsNil)
	c.Check(attempts, check.HasLen, 1)
//.........这里部分代码省略.........