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
}