o11y/pkg/query-service/app/inframetrics/hosts.go

package inframetrics

import (
	"context"
	"errors"
	"fmt"
	"math"
	"sort"
	"strings"
	"time"

	"log/slog"

	"github.com/hanzoai/o11y/pkg/query-service/app/metrics/v4/helpers"
	"github.com/hanzoai/o11y/pkg/query-service/common"
	"github.com/hanzoai/o11y/pkg/query-service/constants"
	"github.com/hanzoai/o11y/pkg/query-service/interfaces"
	"github.com/hanzoai/o11y/pkg/query-service/model"
	v3 "github.com/hanzoai/o11y/pkg/query-service/model/v3"
	"github.com/hanzoai/o11y/pkg/query-service/postprocess"
	"github.com/hanzoai/o11y/pkg/types/ctxtypes"

	"github.com/hanzoai/o11y/pkg/types/instrumentationtypes"
	"github.com/hanzoai/o11y/pkg/valuer"
	"golang.org/x/exp/maps"
	"golang.org/x/exp/slices"
)

type HostsRepo struct {
	reader    interfaces.Reader
	querierV2 interfaces.Querier
}

var (
	// we don't have a way to get the resource attributes from the current time series table
	// but we only want to suggest resource attributes for system metrics,
	// this is a list of attributes that we skip from all labels as they are data point attributes
	// TODO(srikanthccv): remove this once we have a way to get resource attributes

	pointAttrsToIgnore = []string{
		"state",
		"cpu",
		"device",
		"direction",
		"mode",
		"mountpoint",
		"type",
		GetDotMetrics("os_type"),
		GetDotMetrics("process_cgroup"),
		GetDotMetrics("process_command"),
		GetDotMetrics("process_command_line"),
		GetDotMetrics("process_executable_name"),
		GetDotMetrics("process_executable_path"),
		GetDotMetrics("process_owner"),
		GetDotMetrics("process_parent_pid"),
		GetDotMetrics("process_pid"),
	}

	queryNamesForTopHosts = map[string][]string{
		"cpu":    {"A", "B", "F1"},
		"memory": {"C", "D", "F2"},
		"wait":   {"E", "F", "F3"},
		"load15": {"G"},
	}

	// TODO(srikanthccv): remove hardcoded metric name and support keys from any system metric
	metricToUseForHostAttributes = GetDotMetrics("system_cpu_load_average_15m")
	hostNameAttrKey              = GetDotMetrics("host_name")
	agentNameToIgnore            = "k8s-infra-otel-agent"
	hostAttrsToEnrich            = []string{
		GetDotMetrics("os_type"),
	}
	metricNamesForHosts = map[string]string{
		"filesystem": GetDotMetrics("system_filesystem_usage"),
		"cpu":        GetDotMetrics("system_cpu_time"),
		"memory":     GetDotMetrics("system_memory_usage"),
		"load15":     GetDotMetrics("system_cpu_load_average_15m"),
		"wait":       GetDotMetrics("system_cpu_time"),
	}
	uniqueMetricNamesForHosts = []string{
		GetDotMetrics("system_uptime"),
		GetDotMetrics("system_cpu_time"),
		GetDotMetrics("system_cpu_load_average_1m"),
		GetDotMetrics("system_cpu_load_average_5m"),
		GetDotMetrics("system_cpu_load_average_15m"),
		GetDotMetrics("system_memory_usage"),
		GetDotMetrics("system_paging_usage"),
		GetDotMetrics("system_paging_faults"),
		GetDotMetrics("system_paging_operations"),
		GetDotMetrics("system_disk_io"),
		GetDotMetrics("system_disk_operations"),
		GetDotMetrics("system_disk_io_time"),
		GetDotMetrics("system_disk_operation_time"),
		GetDotMetrics("system_disk_merged"),
		GetDotMetrics("system_disk_pending_operations"),
		GetDotMetrics("system_disk_weighted_io_time"),
		GetDotMetrics("system_filesystem_usage"),
		GetDotMetrics("system_filesystem_inodes_usage"),
		GetDotMetrics("system_network_io"),
		GetDotMetrics("system_network_errors"),
		GetDotMetrics("system_network_connections"),
		GetDotMetrics("system_network_dropped"),
		GetDotMetrics("system_network_packets"),
		GetDotMetrics("system_processes_count"),
		GetDotMetrics("system_processes_created"),
		GetDotMetrics("process_cpu_time"),
		GetDotMetrics("process_disk_io"),
		GetDotMetrics("process_memory_usage"),
		GetDotMetrics("process_memory_virtual"),
		GetDotMetrics("nfs_client_net_count"),
		GetDotMetrics("nfs_client_net_tcp_connection_accepted"),
		GetDotMetrics("nfs_client_operation_count"),
		GetDotMetrics("nfs_client_procedure_count"),
		GetDotMetrics("nfs_client_rpc_authrefresh_count"),
		GetDotMetrics("nfs_client_rpc_count"),
		GetDotMetrics("nfs_client_rpc_retransmit_count"),
		GetDotMetrics("nfs_server_fh_stale_count"),
		GetDotMetrics("nfs_server_io"),
		GetDotMetrics("nfs_server_net_count"),
		GetDotMetrics("nfs_server_net_tcp_connection_accepted"),
		GetDotMetrics("nfs_server_operation_count"),
		GetDotMetrics("nfs_server_procedure_count"),
		GetDotMetrics("nfs_server_repcache_requests"),
		GetDotMetrics("nfs_server_rpc_count"),
		GetDotMetrics("nfs_server_thread_count"),
	}
)

func NewHostsRepo(reader interfaces.Reader, querierV2 interfaces.Querier) *HostsRepo {
	return &HostsRepo{reader: reader, querierV2: querierV2}
}

func (h *HostsRepo) GetHostAttributeKeys(ctx context.Context, orgID valuer.UUID, req v3.FilterAttributeKeyRequest) (*v3.FilterAttributeKeyResponse, error) {
	// TODO(srikanthccv): remove hardcoded metric name and support keys from any system metric
	req.DataSource = v3.DataSourceMetrics
	req.AggregateAttribute = metricToUseForHostAttributes
	if req.Limit == 0 {
		req.Limit = 50
	}

	attributeKeysResponse, err := h.reader.GetMetricAttributeKeys(ctx, orgID, &req)
	if err != nil {
		return nil, err
	}

	// TODO(srikanthccv): only return resource attributes when we have a way to
	// distinguish between resource attributes and other attributes.
	filteredKeys := []v3.AttributeKey{}
	for _, key := range attributeKeysResponse.AttributeKeys {
		if slices.Contains(pointAttrsToIgnore, key.Key) {
			continue
		}
		filteredKeys = append(filteredKeys, key)
	}

	return &v3.FilterAttributeKeyResponse{AttributeKeys: filteredKeys}, nil
}

func (h *HostsRepo) GetHostAttributeValues(ctx context.Context, orgID valuer.UUID, req v3.FilterAttributeValueRequest) (*v3.FilterAttributeValueResponse, error) {
	req.DataSource = v3.DataSourceMetrics
	req.AggregateAttribute = metricToUseForHostAttributes
	if req.Limit == 0 {
		req.Limit = 50
	}

	attributeValuesResponse, err := h.reader.GetMetricAttributeValues(ctx, orgID, &req)
	if err != nil {
		return nil, err
	}
	if req.FilterAttributeKey != hostNameAttrKey {
		return attributeValuesResponse, nil
	}
	hostNames := []string{}

	for _, attributeValue := range attributeValuesResponse.StringAttributeValues {
		if strings.Contains(attributeValue, agentNameToIgnore) {
			continue
		}
		hostNames = append(hostNames, attributeValue)
	}

	return &v3.FilterAttributeValueResponse{StringAttributeValues: hostNames}, nil
}

func (h *HostsRepo) getActiveHosts(ctx context.Context) (map[string]bool, error) {
	tenMinAgo := time.Now().Add(-10 * time.Minute).UTC().UnixMilli()
	return h.reader.GetActiveHostsFromMetricMetadata(ctx, uniqueMetricNamesForHosts, hostNameAttrKey, tenMinAgo)
}

func (h *HostsRepo) getMetadataAttributes(ctx context.Context, req model.HostListRequest) (map[string]map[string]string, error) {
	hostAttrs := map[string]map[string]string{}

	for _, key := range hostAttrsToEnrich {
		hasKey := false
		for _, groupByKey := range req.GroupBy {
			if groupByKey.Key == key {
				hasKey = true
				break
			}
		}
		if !hasKey {
			req.GroupBy = append(req.GroupBy, v3.AttributeKey{Key: key})
		}
	}

	mq := v3.BuilderQuery{
		DataSource: v3.DataSourceMetrics,
		AggregateAttribute: v3.AttributeKey{
			Key:      metricToUseForHostAttributes,
			DataType: v3.AttributeKeyDataTypeFloat64,
		},
		Temporality: v3.Unspecified,
		GroupBy:     req.GroupBy,
	}

	query, err := helpers.PrepareTimeseriesFilterQuery(req.Start, req.End, &mq)
	if err != nil {
		return nil, err
	}

	query = localQueryToDistributedQuery(query)

	attrsListResponse, err := h.reader.GetListResultV3(ctx, query)
	if err != nil {
		return nil, err
	}

	for _, row := range attrsListResponse {
		stringData := map[string]string{}
		for key, value := range row.Data {
			if str, ok := value.(string); ok {
				stringData[key] = str
			} else if strPtr, ok := value.(*string); ok {
				stringData[key] = *strPtr
			}
		}

		hostName := stringData[hostNameAttrKey]
		if _, ok := hostAttrs[hostName]; !ok {
			hostAttrs[hostName] = map[string]string{}
		}

		for _, key := range req.GroupBy {
			hostAttrs[hostName][key.Key] = stringData[key.Key]
		}
	}

	return hostAttrs, nil
}

func (h *HostsRepo) getTopHostGroups(ctx context.Context, orgID valuer.UUID, req model.HostListRequest, q *v3.QueryRangeParamsV3) ([]map[string]string, []map[string]string, error) {
	step, timeSeriesTableName, samplesTableName := getParamsForTopHosts(req)

	queryNames := queryNamesForTopHosts[req.OrderBy.ColumnName]
	topHostGroupsQueryRangeParams := &v3.QueryRangeParamsV3{
		Start: req.Start,
		End:   req.End,
		Step:  step,
		CompositeQuery: &v3.CompositeQuery{
			BuilderQueries: map[string]*v3.BuilderQuery{},
			QueryType:      v3.QueryTypeBuilder,
			PanelType:      v3.PanelTypeTable,
		},
	}

	for _, queryName := range queryNames {
		query := q.CompositeQuery.BuilderQueries[queryName].Clone()
		query.StepInterval = step
		query.MetricTableHints = &v3.MetricTableHints{
			TimeSeriesTableName: timeSeriesTableName,
			SamplesTableName:    samplesTableName,
		}
		if req.Filters != nil && len(req.Filters.Items) > 0 {
			query.Filters.Items = append(query.Filters.Items, req.Filters.Items...)
		}
		topHostGroupsQueryRangeParams.CompositeQuery.BuilderQueries[queryName] = query
	}

	ctx = ctxtypes.NewContextWithCommentVals(ctx, map[string]string{
		instrumentationtypes.CodeNamespace:    "inframetrics",
		instrumentationtypes.CodeFunctionName: "getTopHostGroups",
	})
	queryResponse, _, err := h.querierV2.QueryRange(ctx, orgID, topHostGroupsQueryRangeParams)
	if err != nil {
		return nil, nil, err
	}
	formattedResponse, err := postprocess.PostProcessResult(queryResponse, topHostGroupsQueryRangeParams)
	if err != nil {
		return nil, nil, err
	}

	if len(formattedResponse) == 0 || len(formattedResponse[0].Series) == 0 {
		return nil, nil, nil
	}

	if req.OrderBy.Order == v3.DirectionDesc {
		sort.Slice(formattedResponse[0].Series, func(i, j int) bool {
			return formattedResponse[0].Series[i].Points[0].Value > formattedResponse[0].Series[j].Points[0].Value
		})
	} else {
		sort.Slice(formattedResponse[0].Series, func(i, j int) bool {
			return formattedResponse[0].Series[i].Points[0].Value < formattedResponse[0].Series[j].Points[0].Value
		})
	}

	limit := math.Min(float64(req.Offset+req.Limit), float64(len(formattedResponse[0].Series)))

	paginatedTopHostGroupsSeries := formattedResponse[0].Series[req.Offset:int(limit)]

	topHostGroups := []map[string]string{}
	for _, series := range paginatedTopHostGroupsSeries {
		topHostGroups = append(topHostGroups, series.Labels)
	}
	allHostGroups := []map[string]string{}
	for _, series := range formattedResponse[0].Series {
		allHostGroups = append(allHostGroups, series.Labels)
	}

	return topHostGroups, allHostGroups, nil
}

// GetHostMetricsExistenceAndEarliestTime returns (count, minFirstReportedUnixMilli, error) for host metrics
// in distributed_metadata. Uses metricNamesForHosts plus system.filesystem.usage.
func (h *HostsRepo) GetHostMetricsExistenceAndEarliestTime(ctx context.Context, req model.HostListRequest) (uint64, uint64, error) {
	names := []string{}
	for _, metricName := range metricNamesForHosts {
		names = append(names, metricName)
	}

	return h.reader.GetMetricsExistenceAndEarliestTime(ctx, names)
}

func (h *HostsRepo) IsSendingK8SAgentMetrics(ctx context.Context, req model.HostListRequest) ([]string, []string, error) {
	names := []string{}
	for _, metricName := range metricNamesForHosts {
		names = append(names, metricName)
	}
	namesStr := "'" + strings.Join(names, "','") + "'"

	queryForRecentFingerprints := fmt.Sprintf(`
	SELECT DISTINCT fingerprint
	FROM %s.%s
	WHERE metric_name IN (%s)
		AND unix_milli >= toUnixTimestamp(now() - INTERVAL 5 MINUTE) * 1000`,
		constants.O11Y_METRIC_DBNAME, constants.O11Y_SAMPLES_V4_TABLENAME, namesStr)

	query := fmt.Sprintf(`
	SELECT DISTINCT JSONExtractString(labels, '%s') as k8s_cluster_name, JSONExtractString(labels, '%s') as k8s_node_name
	FROM %s.%s
	WHERE metric_name IN (%s)
		AND unix_milli >= toUnixTimestamp(now() - INTERVAL 60 MINUTE) * 1000
		AND JSONExtractString(labels, '%s') LIKE '%%-otel-agent%%'
		AND fingerprint GLOBAL IN (%s)`,
		GetDotMetrics("k8s_cluster_name"), GetDotMetrics("k8s_node_name"),
		constants.O11Y_METRIC_DBNAME, constants.O11Y_TIMESERIES_V4_TABLENAME, namesStr, GetDotMetrics("host_name"), queryForRecentFingerprints)

	result, err := h.reader.GetListResultV3(ctx, query)
	if err != nil {
		return nil, nil, err
	}

	clusterNames := make(map[string]struct{})
	nodeNames := make(map[string]struct{})

	for _, row := range result {
		switch v := row.Data[GetDotMetrics("k8s_cluster_name")].(type) {
		case string:
			clusterNames[v] = struct{}{}
		case *string:
			clusterNames[*v] = struct{}{}
		}
		switch v := row.Data[GetDotMetrics("k8s_node_name")].(type) {
		case string:
			nodeNames[v] = struct{}{}
		case *string:
			nodeNames[*v] = struct{}{}
		}
	}

	return maps.Keys(clusterNames), maps.Keys(nodeNames), nil
}

func (h *HostsRepo) GetHostList(ctx context.Context, orgID valuer.UUID, req model.HostListRequest) (model.HostListResponse, error) {
	resp := model.HostListResponse{}

	if req.Limit == 0 {
		req.Limit = 10
	}

	// default to cpu order by
	if req.OrderBy == nil {
		req.OrderBy = &v3.OrderBy{ColumnName: "cpu", Order: v3.DirectionDesc}
	}

	// default to host name group by
	if len(req.GroupBy) == 0 {
		req.GroupBy = []v3.AttributeKey{{Key: hostNameAttrKey}}
		resp.Type = model.ResponseTypeList
	} else {
		resp.Type = model.ResponseTypeGroupedList
	}

	// don't fail the request if we can't get these values
	if clusterNames, nodeNames, err := h.IsSendingK8SAgentMetrics(ctx, req); err == nil {
		resp.IsSendingK8SAgentMetrics = len(clusterNames) > 0 || len(nodeNames) > 0
		resp.ClusterNames = clusterNames
		resp.NodeNames = nodeNames
	}

	// 1. Check if any host metrics exist and get earliest retention time
	// if no hosts metrics exist, that means we should show the onboarding guide on UI, and return early.
	// 2. If host metrics exist, but req.End is earlier than the earliest time of host metrics as read from
	// metadata table, then we should convey the same to the user and return early
	if count, minFirstReportedUnixMilli, err := h.GetHostMetricsExistenceAndEarliestTime(ctx, req); err == nil {
		if count == 0 {
			resp.SentAnyHostMetricsData = false
			resp.Records = []model.HostListRecord{}
			resp.Total = 0
			return resp, nil
		}
		resp.SentAnyHostMetricsData = true
		if req.End < int64(minFirstReportedUnixMilli) {
			resp.EndTimeBeforeRetention = true
			resp.Records = []model.HostListRecord{}
			resp.Total = 0
			return resp, nil
		}
	}

	step := int64(math.Max(float64(common.MinAllowedStepInterval(req.Start, req.End)), 60))
	if step <= 0 {
		slog.ErrorContext(ctx, "step is less than or equal to 0", "step", step)
		return resp, errors.New("step is less than or equal to 0")
	}

	query := HostsTableListQuery.Clone()

	query.Start = req.Start
	query.End = req.End
	query.Step = step

	for _, query := range query.CompositeQuery.BuilderQueries {
		query.StepInterval = step
		if req.Filters != nil && len(req.Filters.Items) > 0 {
			query.Filters.Items = append(query.Filters.Items, req.Filters.Items...)
		}
		query.GroupBy = req.GroupBy
	}

	hostAttrs, err := h.getMetadataAttributes(ctx, req)
	if err != nil {
		return resp, err
	}

	activeHosts, err := h.getActiveHosts(ctx)
	if err != nil {
		return resp, err
	}

	topHostGroups, allHostGroups, err := h.getTopHostGroups(ctx, orgID, req, query)
	if err != nil {
		return resp, err
	}

	groupFilters := map[string][]string{}
	for _, topHostGroup := range topHostGroups {
		for k, v := range topHostGroup {
			groupFilters[k] = append(groupFilters[k], v)
		}
	}

	for groupKey, groupValues := range groupFilters {
		hasGroupFilter := false
		if req.Filters != nil && len(req.Filters.Items) > 0 {
			for _, filter := range req.Filters.Items {
				if filter.Key.Key == groupKey {
					hasGroupFilter = true
					break
				}
			}
		}
		if !hasGroupFilter {
			for _, query := range query.CompositeQuery.BuilderQueries {
				query.Filters.Items = append(query.Filters.Items, v3.FilterItem{
					Key:      v3.AttributeKey{Key: groupKey},
					Value:    groupValues,
					Operator: v3.FilterOperatorIn,
				})
			}
		}
	}

	ctx = ctxtypes.NewContextWithCommentVals(ctx, map[string]string{
		instrumentationtypes.CodeNamespace:    "inframetrics",
		instrumentationtypes.CodeFunctionName: "GetHostList",
	})
	queryResponse, _, err := h.querierV2.QueryRange(ctx, orgID, query)
	if err != nil {
		return resp, err
	}

	formattedResponse, err := postprocess.PostProcessResult(queryResponse, query)
	if err != nil {
		return resp, err
	}

	records := []model.HostListRecord{}

	for _, result := range formattedResponse {
		for _, row := range result.Table.Rows {
			record := model.HostListRecord{
				CPU:    -1,
				Memory: -1,
				Wait:   -1,
				Load15: -1,
			}

			if hostName, ok := row.Data[hostNameAttrKey].(string); ok {
				record.HostName = hostName
			}

			if cpu, ok := row.Data["F1"].(float64); ok {
				record.CPU = cpu
			}
			if memory, ok := row.Data["F2"].(float64); ok {
				record.Memory = memory
			}
			if wait, ok := row.Data["F3"].(float64); ok {
				record.Wait = wait
			}
			if load15, ok := row.Data["G"].(float64); ok {
				record.Load15 = load15
			}
			record.Meta = map[string]string{}
			if _, ok := hostAttrs[record.HostName]; ok {
				record.Meta = hostAttrs[record.HostName]
			}
			if osType, ok := record.Meta[GetDotMetrics("os_type")]; ok {
				record.OS = osType
			}
			record.Active = activeHosts[record.HostName]
			records = append(records, record)
		}
	}
	resp.Total = len(allHostGroups)
	resp.Records = records
	resp.SortBy(req.OrderBy)

	return resp, nil
}