Golang pkg.NewUnsafeSet函数代码示例

// units returns a map representing the current state of units known by the agent.
func (a *Agent) units() (unitStates, error) {
	launched := pkg.NewUnsafeSet()
	for _, jName := range a.cache.launchedJobs() {
		launched.Add(jName)
	}

	loaded := pkg.NewUnsafeSet()
	for _, jName := range a.cache.loadedJobs() {
		loaded.Add(jName)
	}

	units, err := a.um.Units()
	if err != nil {
		return nil, fmt.Errorf("failed fetching loaded units from UnitManager: %v", err)
	}

	filter := pkg.NewUnsafeSet()
	for _, u := range units {
		filter.Add(u)
	}

	states := make(unitStates)
	for _, uName := range units {
		js := job.JobStateInactive
		if loaded.Contains(uName) {
			js = job.JobStateLoaded
		} else if launched.Contains(uName) {
			js = job.JobStateLaunched
		}
		states[uName] = js
	}

	return states, nil
}

func TestHasMetadata(t *testing.T) {
	testCases := []struct {
		metadata map[string]string
		match    map[string]pkg.Set
		want     bool
	}{
		{
			map[string]string{
				"region": "us-east-1",
			},
			map[string]pkg.Set{
				"region": pkg.NewUnsafeSet("us-east-1"),
			},
			true,
		},
		{
			map[string]string{
				"groups": "ping",
			},
			map[string]pkg.Set{
				"groups": pkg.NewUnsafeSet("ping", "pong"),
			},
			true,
		},
		{
			map[string]string{
				"groups": "ping",
			},
			map[string]pkg.Set{
				"groups": pkg.NewUnsafeSet("pong"),
			},
			false,
		},
		{
			map[string]string{
				"region": "us-east-1",
				"groups": "ping",
			},
			map[string]pkg.Set{
				"region": pkg.NewUnsafeSet("us-east-1"),
				"groups": pkg.NewUnsafeSet("pong"),
			},
			false,
		},
	}

	for i, tt := range testCases {
		ms := &MachineState{Metadata: tt.metadata}
		got := HasMetadata(ms, tt.match)
		if got != tt.want {
			t.Errorf("case %d: HasMetadata returned %t, expected %t", i, got, tt.want)
		}
	}
}

// jobs returns a collection of all Jobs that the Agent has either loaded
// or launched. The Unit, TargetState and TargetMachineID fields of the
// returned *job.Job objects are not properly hydrated.
func (a *Agent) jobs() (map[string]*job.Job, error) {
	launched := pkg.NewUnsafeSet()
	for _, jName := range a.cache.launchedJobs() {
		launched.Add(jName)
	}

	loaded := pkg.NewUnsafeSet()
	for _, jName := range a.cache.loadedJobs() {
		loaded.Add(jName)
	}

	units, err := a.um.Units()
	if err != nil {
		return nil, fmt.Errorf("failed fetching loaded units from UnitManager: %v", err)
	}

	filter := pkg.NewUnsafeSet()
	for _, u := range units {
		filter.Add(u)
	}
	states, err := a.um.GetUnitStates(filter)
	if err != nil {
		return nil, fmt.Errorf("failed fetching unit states: %v", err)
	}

	jobs := make(map[string]*job.Job)
	for _, uName := range units {
		jobs[uName] = &job.Job{
			Name:      uName,
			UnitState: states[uName],
			State:     nil,

			// The following fields are not properly populated
			// and should not be used in the calling code
			Unit:            unit.Unit{},
			TargetState:     job.JobState(""),
			TargetMachineID: "",
		}

		js := job.JobStateInactive
		if loaded.Contains(uName) {
			js = job.JobStateLoaded
		} else if launched.Contains(uName) {
			js = job.JobStateLaunched
		}
		jobs[uName].State = &js
	}

	return jobs, nil
}

func TestFakeUnitManagerGetUnitStates(t *testing.T) {
	fum := NewFakeUnitManager()

	err := fum.Load("hello.service", Unit{})
	if err != nil {
		t.Fatalf("Expected no error from Load(), got %v", err)
	}

	states, err := fum.GetUnitStates(pkg.NewUnsafeSet("hello.service", "goodbye.service"))
	if err != nil {
		t.Fatalf("Failed calling GetUnitStates: %v", err)
	}

	expectStates := map[string]*UnitState{
		"hello.service": &UnitState{
			LoadState:   "loaded",
			ActiveState: "active",
			SubState:    "running",
		},
	}

	if !reflect.DeepEqual(expectStates, states) {
		t.Fatalf("Received unexpected collection of UnitStates: %#v\nExpected: %#v", states, expectStates)
	}
}

// calculateTaskChainsForJobs compares the desired and current state of an Agent.
// The generated taskChains represent what should be done to make the desired
// state match the current state.
func (ar *AgentReconciler) calculateTaskChainsForJobs(dState, cState *AgentState) <-chan taskChain {
	tcChan := make(chan taskChain)
	go func() {
		jobs := pkg.NewUnsafeSet()
		for cName := range cState.Jobs {
			jobs.Add(cName)
		}

		for dName := range dState.Jobs {
			jobs.Add(dName)
		}

		for _, name := range jobs.Values() {
			tc := ar.calculateTaskChainForJob(dState, cState, name)
			if tc == nil {
				continue
			}
			tcChan <- *tc
		}

		close(tcChan)
	}()

	return tcChan
}

// RequiredTargetMetadata return all machine-related metadata from a Job's
// requirements. Valid metadata fields are strings of the form `key=value`,
// where both key and value are not the empty string.
func (j *Job) RequiredTargetMetadata() map[string]pkg.Set {
	metadata := make(map[string]pkg.Set)

	for _, key := range []string{
		deprecatedXConditionPrefix + fleetMachineMetadata,
		fleetMachineMetadata,
	} {
		for _, valuePair := range j.requirements()[key] {
			s := strings.Split(valuePair, "=")

			if len(s) != 2 {
				continue
			}

			if len(s[0]) == 0 || len(s[1]) == 0 {
				continue
			}

			if _, ok := metadata[s[0]]; !ok {
				metadata[s[0]] = pkg.NewUnsafeSet()
			}
			metadata[s[0]].Add(s[1])
		}
	}

	return metadata
}

func newClusterState(jobs []job.Job, unresolved []job.JobOffer, machines []machine.MachineState) *clusterState {
	oSet := pkg.NewUnsafeSet()
	for _, offer := range unresolved {
		oSet.Add(offer.Job.Name)
	}

	mSet := pkg.NewUnsafeSet()
	for _, m := range machines {
		mSet.Add(m.ID)
	}

	return &clusterState{
		jobs:     jobs,
		offers:   oSet,
		machines: mSet,
	}
}

func (f *FakeRegistry) SubmitJobBid(jName, machID string) {
	f.Lock()
	defer f.Unlock()

	_, ok := f.bids[jName]
	if !ok {
		f.bids[jName] = pkg.NewUnsafeSet()
	}
	f.bids[jName].Add(machID)
}

// ValidateOptions ensures that a set of UnitOptions is valid; if not, an error
// is returned detailing the issue encountered.  If there are several problems
// with a set of options, only the first is returned.
func ValidateOptions(opts []*schema.UnitOption) error {
	uf := schema.MapSchemaUnitOptionsToUnitFile(opts)
	j := &job.Job{
		Unit: *uf,
	}
	conflicts := pkg.NewUnsafeSet(j.Conflicts()...)
	peers := pkg.NewUnsafeSet(j.Peers()...)
	for _, peer := range peers.Values() {
		for _, conflict := range conflicts.Values() {
			matched, _ := path.Match(conflict, peer)
			if matched {
				return fmt.Errorf("unresolvable requirements: peer %q matches conflict %q", peer, conflict)
			}
		}
	}
	hasPeers := peers.Length() != 0
	hasConflicts := conflicts.Length() != 0
	_, hasReqTarget := j.RequiredTarget()
	u := &job.Unit{
		Unit: *uf,
	}
	isGlobal := u.IsGlobal()

	switch {
	case hasReqTarget && hasPeers:
		return errors.New("MachineID cannot be used with Peers")
	case hasReqTarget && hasConflicts:
		return errors.New("MachineID cannot be used with Conflicts")
	case hasReqTarget && isGlobal:
		return errors.New("MachineID cannot be used with Global")
	case isGlobal && hasPeers:
		return errors.New("Global cannot be used with Peers")
	case isGlobal && hasConflicts:
		return errors.New("Global cannot be used with Conflicts")
	}

	return nil
}

// calculateTasksForUnits compares the desired and current state of an Agent.
// The generated tasks represent what, in order, should be done to make the
//  desired state match the current state.
func (ar *AgentReconciler) calculateTasksForUnits(dState *AgentState, cState unitStates) []task {
	jobs := pkg.NewUnsafeSet()
	for cName := range cState {
		jobs.Add(cName)
	}

	for dName := range dState.Units {
		jobs.Add(dName)
	}

	var tasks []task
	for _, name := range jobs.Values() {
		tasks = append(tasks, ar.calculateTasksForUnit(dState, cState, name)...)
	}

	sort.Sort(sortableTasks(tasks))
	return tasks
}

func TestTaskManagerUnitSerialization(t *testing.T) {
	result := make(chan error)
	testMapper := func(task, *job.Unit, *Agent) (exec func() error, err error) {
		exec = func() error {
			return <-result
		}
		return
	}

	tm := taskManager{
		processing: pkg.NewUnsafeSet(),
		mapper:     testMapper,
	}

	reschan, err := tm.Do(taskChain{unit: &job.Unit{Name: "foo"}, tasks: []task{task{typ: "test"}}}, nil)
	if err != nil {
		t.Fatalf("unable to start first task: %v", err)
	}

	// the first task should block the second, as it is still in progress
	_, err = tm.Do(taskChain{unit: &job.Unit{Name: "foo"}, tasks: []task{task{typ: "test"}}}, nil)
	if err == nil {
		t.Fatalf("expected error from attempt to start second task, got nil")
	}

	result <- nil

	select {
	case res := <-reschan:
		if res.err != nil {
			t.Errorf("received unexpected error from first task: %v", err)
		}
	default:
		t.Errorf("expected reschan to be ready to receive")
	}

	// since the first task completed, this third task can start
	_, err = tm.Do(taskChain{unit: &job.Unit{Name: "foo"}, tasks: []task{task{typ: "test"}}}, nil)
	if err != nil {
		t.Fatalf("unable to start third task: %v", err)
	}

	close(result)
}

func TestTaskManagerTwoInFlight(t *testing.T) {
	result := make(chan error)
	testMapper := func(task, *job.Unit, *Agent) (exec func() error, err error) {
		exec = func() error {
			return <-result
		}
		return
	}

	tm := taskManager{
		processing: pkg.NewUnsafeSet(),
		mapper:     testMapper,
	}

	errchan1, err := tm.Do(taskChain{unit: &job.Unit{Name: "foo"}, tasks: []task{task{typ: "test"}}}, nil)
	if err != nil {
		t.Fatalf("unable to start task: %v", err)
	}

	errchan2, err := tm.Do(taskChain{unit: &job.Unit{Name: "bar"}, tasks: []task{task{typ: "test"}}}, nil)
	if err != nil {
		t.Fatalf("unable to start task: %v", err)
	}

	close(result)

	go func() {
		<-time.After(time.Second)
		t.Fatalf("expected errchans to be ready to receive within 1s")
	}()

	res := <-errchan1
	if res.err != nil {
		t.Fatalf("received unexpected error from task one: %v", res.err)
	}

	res = <-errchan2
	if res.err != nil {
		t.Fatalf("received unexpected error from task two: %v", res.err)
	}
}

// calculateTasksForJobs compares the desired and current state of an Agent.
// The generateed tasks represent what should be done to make the desired
// state match the current state.
func (ar *AgentReconciler) calculateTasksForJobs(ms *machine.MachineState, dState, cState *agentState) <-chan *task {
	taskchan := make(chan *task)
	go func() {
		jobs := pkg.NewUnsafeSet()
		for cName := range cState.jobs {
			jobs.Add(cName)
		}

		for dName := range dState.jobs {
			jobs.Add(dName)
		}

		for _, name := range jobs.Values() {
			ar.calculateTasksForJob(ms, dState, cState, name, taskchan)
		}

		close(taskchan)
	}()

	return taskchan
}

// Bids returns a list of machine IDs that have bid for the referenced Job
func (r *EtcdRegistry) Bids(jName string) (bids pkg.Set, err error) {
	bids = pkg.NewUnsafeSet()

	req := etcd.Get{
		Key:       path.Join(r.keyPrefix, offerPrefix, jName, "bids"),
		Recursive: true,
	}

	var resp *etcd.Result
	resp, err = r.etcd.Do(&req)
	if err != nil {
		if isKeyNotFound(err) {
			err = nil
		}
		return
	}

	for _, node := range resp.Node.Nodes {
		machID := path.Base(node.Key)
		bids.Add(machID)
	}

	return
}

func TestCalculateTasksForOffer(t *testing.T) {
	tests := []struct {
		mState *machine.MachineState
		dState *agentState
		job    *job.Job
		bids   pkg.Set

		tasks []task
	}{
		// no bid submitted yet and able to run
		{
			mState: &machine.MachineState{ID: "XXX"},
			dState: newAgentState(),
			job: &job.Job{
				Name:        "foo.service",
				TargetState: jsLaunched,
				Unit:        fleetUnit(t),
			},
			bids: pkg.NewUnsafeSet(),
			tasks: []task{
				task{
					Type: taskTypeSubmitBid,
					Job: &job.Job{
						Name:        "foo.service",
						TargetState: jsLaunched,
						Unit:        fleetUnit(t),
					},
					Reason: taskReasonAbleToResolveOffer,
				},
			},
		},

		// no bid submitted but unable to run
		{
			mState: &machine.MachineState{ID: "XXX"},
			dState: newAgentState(),
			job: &job.Job{
				Name:        "foo.service",
				TargetState: jsLaunched,
				Unit:        fleetUnit(t, "X-ConditionMachineID=YYY"),
			},
			bids:  pkg.NewUnsafeSet(),
			tasks: []task{},
		},

		// bid already submitted
		{
			mState: &machine.MachineState{ID: "XXX"},
			dState: newAgentState(),
			job: &job.Job{
				TargetState: jsLaunched,
				Unit:        fleetUnit(t),
			},
			bids:  pkg.NewUnsafeSet("XXX"),
			tasks: []task{},
		},
	}

	for i, tt := range tests {
		ar := NewReconciler(registry.NewFakeRegistry(), nil)
		taskchan := make(chan *task)
		tasks := []task{}
		go func() {
			ar.calculateTasksForOffer(tt.dState, tt.mState, tt.job, tt.bids, taskchan)
			close(taskchan)
		}()

		for t := range taskchan {
			tasks = append(tasks, *t)
		}

		if !reflect.DeepEqual(tt.tasks, tasks) {
			t.Errorf("case %d: calculated incorrect list of tasks\nexpected=%v\nreceived=%v\n", i, tt.tasks, tasks)
		}
	}
}