func()

in pkg/hostmgr/hostsvc/handler.go [69:155]
78 lines of code
25 McCabe index (conditional complexity)

func (m *serviceHandler) QueryHosts(
	ctx context.Context,
	request *host_svc.QueryHostsRequest,
) (*host_svc.QueryHostsResponse, error) {
	m.metrics.QueryHostsAPI.Inc(1)

	// Add request.HostStates to a set to remove duplicates
	hostStateSet := stringset.New()
	for _, state := range request.GetHostStates() {
		hostStateSet.Add(state.String())
	}

	if request.HostStates == nil || len(request.HostStates) == 0 {
		for _, state := range hpb.HostState_name {
			hostStateSet.Add(state)
		}
	}

	var hostInfos []*hpb.HostInfo
	drainingHostsInfo, err := m.drainer.GetAllDrainingHostInfos()
	if err != nil {
		return nil, err
	}
	drainedHostsInfo, err := m.drainer.GetAllDrainedHostInfos()
	if err != nil {
		return nil, err
	}
	downHostsInfo, err := m.drainer.GetAllDownHostInfos()
	if err != nil {
		return nil, err
	}
	for _, hostState := range hostStateSet.ToSlice() {
		switch hostState {
		case hpb.HostState_HOST_STATE_UP.String():
			upHosts, err := host.BuildHostInfoForRegisteredAgents()
			if err != nil {
				m.metrics.QueryHostsFail.Inc(1)
				return nil, yarpcerrors.InternalErrorf(err.Error())
			}
			// Remove draining / drained / down hosts from the result.
			// This is needed because AgentMap is updated every 15s
			// and might not have the up to date information.
			for _, hostInfo := range drainingHostsInfo {
				delete(upHosts, hostInfo.GetHostname())
			}
			for _, hostInfo := range drainedHostsInfo {
				delete(upHosts, hostInfo.GetHostname())
			}
			for _, hostInfo := range downHostsInfo {
				delete(upHosts, hostInfo.GetHostname())
			}

			for _, hostInfo := range upHosts {
				hostInfos = append(hostInfos, hostInfo)
			}
		case hpb.HostState_HOST_STATE_DRAINING.String():
			for _, hostInfo := range drainingHostsInfo {
				hostInfos = append(hostInfos, hostInfo)
			}
		case hpb.HostState_HOST_STATE_DRAINED.String():
			for _, hostInfo := range drainedHostsInfo {
				hostInfos = append(hostInfos, hostInfo)
			}
		case hpb.HostState_HOST_STATE_DOWN.String():
			for _, hostInfo := range downHostsInfo {
				hostInfos = append(hostInfos, hostInfo)
			}
		}
	}
	if m.hostPoolManager != nil {
		for _, h := range hostInfos {
			p, err := m.hostPoolManager.GetPoolByHostname(h.GetHostname())
			if err == nil {
				h.CurrentPool = p.ID()
			}
			d, err := m.hostPoolManager.GetDesiredPool(h.GetHostname())
			if err == nil {
				h.DesiredPool = d
			}
		}
	}

	m.metrics.QueryHostsSuccess.Inc(1)
	return &host_svc.QueryHostsResponse{
		HostInfos: hostInfos,
	}, nil
}