func newFakeRegistryForCommands(unitPrefix string, unitCount int, template bool) client.API {
// clear machineStates for every invocation
machineStates = nil
machines := []machine.MachineState{
newMachineState("c31e44e1-f858-436e-933e-59c642517860", "1.2.3.4", map[string]string{"ping": "pong"}),
newMachineState("595989bb-cbb7-49ce-8726-722d6e157b4e", "5.6.7.8", map[string]string{"foo": "bar"}),
}
jobs := make([]job.Job, 0)
appendJobsForTests(&jobs, machines[0], unitPrefix, unitCount, template)
appendJobsForTests(&jobs, machines[1], unitPrefix, unitCount, template)
states := make([]unit.UnitState, 0)
if template {
state := unit.UnitState{
UnitName: fmt.Sprintf("%[email protected]", unitPrefix),
LoadState: "loaded",
ActiveState: "inactive",
SubState: "dead",
MachineID: machines[0].ID,
}
states = append(states, state)
state.MachineID = machines[1].ID
states = append(states, state)
} else {
for i := 1; i <= unitCount; i++ {
state := unit.UnitState{
UnitName: fmt.Sprintf("%s%d.service", unitPrefix, i),
LoadState: "loaded",
ActiveState: "active",
SubState: "listening",
MachineID: machines[0].ID,
}
states = append(states, state)
}
for i := 1; i <= unitCount; i++ {
state := unit.UnitState{
UnitName: fmt.Sprintf("%s%d.service", unitPrefix, i),
LoadState: "loaded",
ActiveState: "inactive",
SubState: "dead",
MachineID: machines[1].ID,
}
states = append(states, state)
}
}
reg := registry.NewFakeRegistry()
reg.SetMachines(machines)
reg.SetUnitStates(states)
reg.SetJobs(jobs)
return &client.RegistryClient{Registry: reg}
}
func (r *RPCRegistry) SaveUnitState(unitName string, unitState *unit.UnitState, ttl time.Duration) {
if DebugRPCRegistry {
defer debug.Exit_(debug.Enter_(unitName, unitState))
}
if unitState.UnitName == "" {
unitState.UnitName = unitName
}
r.getClient().SaveUnitState(r.ctx(), &pb.SaveUnitStateRequest{
Name: unitName,
State: unitState.ToPB(),
TTL: int32(ttl.Seconds()),
})
}
func modelToUnitState(usm *unitStateModel) *unit.UnitState {
if usm == nil {
return nil
}
us := unit.UnitState{
LoadState: usm.LoadState,
ActiveState: usm.ActiveState,
SubState: usm.SubState,
UnitHash: usm.UnitHash,
}
if usm.MachineState != nil {
us.MachineID = usm.MachineState.ID
}
return &us
}
// ReportUnitState attaches the current state of the Agent's Machine to the given
// unit.UnitState object, then persists that state in the Registry
func (a *Agent) ReportUnitState(jobName string, us *unit.UnitState) {
if us == nil {
log.V(1).Infof("Job(%s): purging UnitState from Registry", jobName)
err := a.registry.RemoveUnitState(jobName)
if err != nil {
log.Errorf("Failed to remove UnitState for job %s from Registry: %s", jobName, err.Error())
}
} else {
ms := a.Machine.State()
us.MachineState = &ms
log.V(1).Infof("Job(%s): pushing UnitState (loadState=%s, activeState=%s, subState=%s) to Registry", jobName, us.LoadState, us.ActiveState, us.SubState)
a.registry.SaveUnitState(jobName, us)
}
}