func (s *rpcserver) GetUnitStates(ctx context.Context, filter *pb.UnitStateFilter) (*pb.UnitStates, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_())
}
states := make([]*pb.UnitState, 0)
states = append(states, s.localRegistry.UnitStates()...)
if s.hasNonGRPCAgents {
log.Debug("Merging etcd with inmemory unit states in GetUnitStates()")
etcdUnitStates, err := s.etcdRegistry.UnitStates()
if err != nil {
return nil, err
}
unitStateNames := make(map[string]string, len(states))
for _, state := range states {
unitStateNames[state.Name] = state.MachineID
}
for _, state := range etcdUnitStates {
machId, ok := unitStateNames[state.UnitName]
if !ok || (ok && machId != state.MachineID) {
states = append(states, state.ToPB())
}
}
}
return &pb.UnitStates{states}, nil
}
func (s *rpcserver) GetUnits(ctx context.Context, filter *pb.UnitFilter) (*pb.Units, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_())
}
units := make([]pb.Unit, 0)
units = append(units, s.localRegistry.Units()...)
// Check if there are etcd fleet-based agents in the cluster to share the state
if s.hasNonGRPCAgents {
log.Debug("Merging etcd with inmemory units in GetUnits()")
etcdUnits, err := s.etcdRegistry.Units()
if err != nil {
return nil, err
}
unitNames := make(map[string]struct{}, len(units))
for _, unit := range units {
unitNames[unit.Name] = struct{}{}
}
for _, unit := range etcdUnits {
if _, ok := unitNames[unit.Name]; !ok {
units = append(units, unit.ToPB())
}
}
}
return &pb.Units{Units: units}, nil
}
func (r *inmemoryRegistry) DestroyUnit(name string) bool {
if DebugInmemoryRegistry {
defer debug.Exit_(debug.Enter_(name))
}
r.mu.Lock()
defer r.mu.Unlock()
r.unitStatesMu.Lock()
defer r.unitStatesMu.Unlock()
r.heartbeatsMu.Lock()
defer r.heartbeatsMu.Unlock()
deleted := false
if _, exists := r.unitsCache[name]; exists {
delete(r.unitsCache, name)
deleted = true
}
if _, exists := r.scheduledUnits[name]; exists {
delete(r.scheduledUnits, name)
}
if _, exists := r.unitHeartbeats[name]; exists {
delete(r.unitHeartbeats, name)
}
if _, exists := r.unitStates[name]; exists {
delete(r.unitStates, name)
}
return deleted
}
func (r *RPCRegistry) ClearUnitHeartbeat(unitName string) {
if DebugRPCRegistry {
defer debug.Exit_(debug.Enter_(unitName))
}
r.getClient().ClearUnitHeartbeat(r.ctx(), &pb.UnitName{unitName})
}
func (s *rpcserver) ClearUnitHeartbeat(ctx context.Context, name *pb.UnitName) (*pb.GenericReply, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(name.Name))
}
s.localRegistry.ClearUnitHeartbeat(name.Name)
return &pb.GenericReply{}, nil
}
func (s *rpcserver) UnitHeartbeat(ctx context.Context, heartbeat *pb.Heartbeat) (*pb.GenericReply, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(heartbeat))
}
s.localRegistry.UnitHeartbeat(heartbeat.Name, heartbeat.MachineID, time.Duration(heartbeat.TTL)*time.Second)
return &pb.GenericReply{}, nil
}
func (r *RPCRegistry) DestroyUnit(unitName string) error {
if DebugRPCRegistry {
defer debug.Exit_(debug.Enter_(unitName))
}
_, err := r.getClient().DestroyUnit(r.ctx(), &pb.UnitName{unitName})
return err
}
func (r *inmemoryRegistry) UnitState(name string) *pb.UnitState {
if DebugInmemoryRegistry {
defer debug.Exit_(debug.Enter_())
}
r.unitStatesMu.Lock()
defer r.unitStatesMu.Unlock()
return r.stateByMUSKey(name)
}
func (r *inmemoryRegistry) CreateUnit(u *pb.Unit) {
if DebugInmemoryRegistry {
defer debug.Exit_(debug.Enter_(u))
}
r.mu.Lock()
defer r.mu.Unlock()
r.unitsCache[u.Name] = *u
}
func (s *rpcserver) GetScheduledUnits(ctx context.Context, unitFilter *pb.UnitFilter) (*pb.ScheduledUnits, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_())
}
units, err := s.localRegistry.Schedule()
return &pb.ScheduledUnits{Units: units}, err
}
func (r *inmemoryRegistry) UnscheduleUnit(unitName, machineid string) {
if DebugInmemoryRegistry {
defer debug.Exit_(debug.Enter_(unitName, machineid))
}
r.mu.Lock()
defer r.mu.Unlock()
delete(r.scheduledUnits, unitName)
}
func (s *rpcserver) GetScheduledUnit(ctx context.Context, name *pb.UnitName) (*pb.MaybeScheduledUnit, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(name.Name))
}
su := s.localRegistry.ScheduledUnit(name.Name)
return &pb.MaybeScheduledUnit{IsScheduled: &pb.MaybeScheduledUnit_Unit{Unit: su}}, nil
}
func (s *rpcserver) GetUnitState(ctx context.Context, name *pb.UnitName) (*pb.UnitState, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(name.Name))
}
state := s.localRegistry.UnitState(name.Name)
return state, nil
}
func (r *RPCRegistry) CreateUnit(j *job.Unit) error {
if DebugRPCRegistry {
defer debug.Exit_(debug.Enter_(j.Name))
}
un := j.ToPB()
_, err := r.getClient().CreateUnit(r.ctx(), &un)
return err
}
func (r *inmemoryRegistry) Unit(name string) (pb.Unit, bool) {
if DebugInmemoryRegistry {
defer debug.Exit_(debug.Enter_(name))
}
r.mu.RLock()
defer r.mu.RUnlock()
u, exists := r.unitsCache[name]
return u, exists
}
func (s *rpcserver) UnscheduleUnit(ctx context.Context, unit *pb.UnscheduleUnitRequest) (*pb.GenericReply, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(unit.Name, unit.MachineID))
}
err := s.etcdRegistry.UnscheduleUnit(unit.Name, unit.MachineID)
if err == nil {
s.localRegistry.UnscheduleUnit(unit.Name, unit.MachineID)
}
return &pb.GenericReply{}, err
}
func (r *inmemoryRegistry) ClearUnitHeartbeat(name string) {
if DebugInmemoryRegistry {
defer debug.Exit_(debug.Enter_(name))
}
r.heartbeatsMu.Lock()
defer r.heartbeatsMu.Unlock()
if _, exists := r.unitHeartbeats[name]; exists {
delete(r.unitHeartbeats, name)
}
}
func (s *rpcserver) CreateUnit(ctx context.Context, u *pb.Unit) (*pb.GenericReply, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(u.Name))
}
err := s.etcdRegistry.CreateUnit(rpcUnitToJobUnit(u))
if err == nil {
s.localRegistry.CreateUnit(u)
}
return &pb.GenericReply{}, err
}
func (s *rpcserver) DestroyUnit(ctx context.Context, name *pb.UnitName) (*pb.GenericReply, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(name.Name))
}
err := s.etcdRegistry.DestroyUnit(name.Name)
if err == nil {
s.localRegistry.DestroyUnit(name.Name)
}
return &pb.GenericReply{}, err
}
func (r *RPCRegistry) UnscheduleUnit(unitName, machID string) error {
if DebugRPCRegistry {
defer debug.Exit_(debug.Enter_(unitName, machID))
}
_, err := r.getClient().UnscheduleUnit(r.ctx(), &pb.UnscheduleUnitRequest{
Name: unitName,
MachineID: machID,
})
return err
}
func (r *RPCRegistry) SetUnitTargetState(unitName string, state job.JobState) error {
if DebugRPCRegistry {
defer debug.Exit_(debug.Enter_(unitName, state))
}
_, err := r.getClient().SetUnitTargetState(r.ctx(), &pb.ScheduledUnit{
Name: unitName,
CurrentState: state.ToPB(),
})
return err
}
func (r *inmemoryRegistry) RemoveUnitState(unitName string) {
if DebugInmemoryRegistry {
defer debug.Exit_(debug.Enter_(unitName))
}
r.unitStatesMu.Lock()
defer r.unitStatesMu.Unlock()
if _, exists := r.unitStates[unitName]; exists {
delete(r.unitStates, unitName)
}
}
func (s *rpcserver) SetUnitTargetState(ctx context.Context, unit *pb.ScheduledUnit) (*pb.GenericReply, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(unit.Name, unit.CurrentState))
}
err := s.etcdRegistry.SetUnitTargetState(unit.Name, rpcUnitStateToJobState(unit.CurrentState))
if err == nil {
if s.localRegistry.SetUnitTargetState(unit.Name, unit.CurrentState) {
}
}
return &pb.GenericReply{}, err
}
func (r *RPCRegistry) UnitHeartbeat(unitName, machID string, ttl time.Duration) error {
if DebugRPCRegistry {
defer debug.Exit_(debug.Enter_(unitName, machID))
}
_, err := r.getClient().UnitHeartbeat(r.ctx(), &pb.Heartbeat{
Name: unitName,
MachineID: machID,
TTL: int32(ttl.Seconds()),
})
return err
}
func (s *rpcserver) GetUnit(ctx context.Context, name *pb.UnitName) (*pb.MaybeUnit, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(name.Name))
}
unit, exists := s.localRegistry.Unit(name.Name)
if exists {
return &pb.MaybeUnit{HasUnit: &pb.MaybeUnit_Unit{Unit: &unit}}, nil
}
return &pb.MaybeUnit{HasUnit: &pb.MaybeUnit_Notfound{Notfound: &pb.NotFound{}}}, nil
}
func (s *rpcserver) RemoveUnitState(ctx context.Context, name *pb.UnitName) (*pb.GenericReply, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(name.Name))
}
// Check if there are etcd fleet-based agents in the cluster to share the state
if s.hasNonGRPCAgents {
s.etcdRegistry.RemoveUnitState(name.Name)
}
s.localRegistry.RemoveUnitState(name.Name)
return &pb.GenericReply{}, nil
}
func (r *inmemoryRegistry) UnitHeartbeat(unitName, machineid string, ttl time.Duration) {
if DebugInmemoryRegistry {
defer debug.Exit_(debug.Enter_(unitName, machineid, ttl))
}
r.heartbeatsMu.Lock()
defer r.heartbeatsMu.Unlock()
if _, exists := r.unitHeartbeats[unitName]; exists {
r.unitHeartbeats[unitName][machineid] = time.Now().Add(ttl)
} else {
r.unitHeartbeats[unitName] = map[string]time.Time{machineid: time.Now().Add(ttl)}
}
}
func (r *inmemoryRegistry) ScheduleUnit(unitName, machineid string) {
if DebugInmemoryRegistry {
defer debug.Exit_(debug.Enter_(unitName, machineid))
}
r.mu.Lock()
defer r.mu.Unlock()
r.scheduledUnits[unitName] = pb.ScheduledUnit{
Name: unitName,
CurrentState: pb.TargetState_INACTIVE,
MachineID: machineid,
}
}
func (s *rpcserver) SaveUnitState(ctx context.Context, req *pb.SaveUnitStateRequest) (*pb.GenericReply, error) {
if debugRPCServer {
defer debug.Exit_(debug.Enter_(req))
}
// Check if there are etcd fleet-based agents in the cluster to share the state
if s.hasNonGRPCAgents {
unitState := rpcUnitStateToExtUnitState(req.State)
s.etcdRegistry.SaveUnitState(req.Name, unitState, time.Duration(req.TTL)*time.Second)
}
s.localRegistry.SaveUnitState(req.Name, req.State, time.Duration(req.TTL)*time.Second)
return &pb.GenericReply{}, nil
}
func (r *inmemoryRegistry) SetUnitTargetState(unitName string, targetState pb.TargetState) bool {
if DebugInmemoryRegistry {
defer debug.Exit_(debug.Enter_(unitName, targetState))
}
r.mu.Lock()
defer r.mu.Unlock()
if u, exists := r.unitsCache[unitName]; exists {
u.DesiredState = targetState
r.unitsCache[unitName] = u
return true
}
return false
}