Thanos源码专题【左扬精讲】——Thanos API 组件（release-0.26）源码阅读和分析（详解 pkg\api\query\grpc.go）

Thanos API 组件（release-0.26）源码阅读和分析（详解 pkg\api\query\grpc.go）

https://github.com/thanos-io/thanos/blob/v0.26.0/pkg/api/query/grpc.go

// Copyright (c) The Thanos Authors.
// Licensed under the Apache License 2.0.

package v1

import (
	"context"
	"time"

	"github.com/prometheus/prometheus/promql"
	"github.com/thanos-io/thanos/pkg/api/query/querypb"
	"github.com/thanos-io/thanos/pkg/query"
	"github.com/thanos-io/thanos/pkg/store/labelpb"
	"github.com/thanos-io/thanos/pkg/store/storepb/prompb"
	"google.golang.org/grpc"
)

type GRPCAPI struct {
	now                         func() time.Time           // now 是一个函数，用于获取当前时间。
	queryableCreate             query.QueryableCreator     // queryableCreate 是一个函数，用于创建查询对象。
	queryEngine                 func(int64) *promql.Engine // queryEngine 是一个函数，用于创建查询引擎。
	defaultMaxResolutionSeconds time.Duration              // defaultMaxResolutionSeconds 是一个时间间隔，用于设置默认的最大分辨率秒数。
}

// NewGRPCAPI 创建一个新的GRPCAPI实例
// now: 当前时间函数，返回当前时间的 time.Time 对象
// creator: 可查询的创建者，实现了 query.QueryableCreator 接口
// queryEngine: 查询引擎函数，接受一个 int64 类型的参数并返回一个指向 promql.Engine 的指针
// defaultMaxResolutionSeconds: 默认的最大解析秒数，类型为 time.Duration
// 返回值：一个指向 GRPCAPI 结构体的指针
func NewGRPCAPI(now func() time.Time, creator query.QueryableCreator, queryEngine func(int64) *promql.Engine, defaultMaxResolutionSeconds time.Duration) *GRPCAPI {
	// 初始化GRPCAPI结构体
	return &GRPCAPI{
		// 当前时间函数
		now: now,
		// 可查询的创建者
		queryableCreate: creator,
		// 查询引擎函数
		queryEngine: queryEngine,
		// 默认的最大解析秒数
		defaultMaxResolutionSeconds: defaultMaxResolutionSeconds,
	}
}

// RegisterQueryServer 函数接收一个 querypb.QueryServer 类型的参数，并返回一个函数。
// 该返回的函数接收一个 *grpc.Server 类型的参数，用于在 gRPC 服务器上注册 QueryServer。
//
// 参数:
//
//	queryServer: querypb.QueryServer 类型的参数，需要被注册的查询服务器。
//
// 返回值:
//
//	返回一个函数，该函数接收一个 *grpc.Server 类型的参数，并在 gRPC 服务器上注册 QueryServer。
func RegisterQueryServer(queryServer querypb.QueryServer) func(*grpc.Server) {
	// 返回一个函数，该函数接收一个 *grpc.Server 类型的参数
	return func(s *grpc.Server) {
		// 在 gRPC 服务器上注册 QueryServer
		// querypb.RegisterQueryServer 是 gRPC 提供的注册函数
		querypb.RegisterQueryServer(s, queryServer)
	}
}

// Query 函数是 GRPCAPI 结构体中的一个方法，用于处理查询请求。
// 它接收一个 *querypb.QueryRequest 类型的请求和一个 querypb.Query_QueryServer 类型的服务器接口，
// 返回一个 error 类型的错误值。
//
// 参数：
//
//	request: *querypb.QueryRequest，包含查询请求信息的结构体指针。
//	server: querypb.Query_QueryServer，用于与客户端通信的服务器接口。
//
// 返回值：
//
//	返回一个 error 类型的错误值，如果查询过程中出现错误则返回该错误，否则返回 nil。
func (g *GRPCAPI) Query(request *querypb.QueryRequest, server querypb.Query_QueryServer) error {
	// 获取上下文环境
	ctx := context.Background()
	var ts time.Time
	// 判断请求中的时间戳是否为0
	if request.TimeSeconds == 0 {
		// 如果为0，则使用当前时间
		ts = g.now()
	} else {
		// 否则，将请求中的时间戳转换为time.Time类型
		ts = time.Unix(request.TimeSeconds, 0)
	}

	// 判断请求中的超时时间是否不为0
	if request.TimeoutSeconds != 0 {
		var cancel context.CancelFunc
		// 将请求中的超时时间转换为time.Duration类型
		timeout := time.Duration(request.TimeoutSeconds) * time.Second
		// 创建一个带超时的上下文环境
		ctx, cancel = context.WithTimeout(ctx, timeout)
		defer cancel() // 确保在函数结束时调用cancel函数
	}

	// 获取请求中的最大分辨率
	maxResolution := request.MaxResolutionSeconds
	// 判断请求中的最大分辨率是否为0
	if request.MaxResolutionSeconds == 0 {
		// 如果为0，则使用默认的最大分辨率
		maxResolution = g.defaultMaxResolutionSeconds.Milliseconds() / 1000
	}

	// 将请求中的StoreMatchers转换为LabelMatchers
	storeMatchers, err := querypb.StoreMatchersToLabelMatchers(request.StoreMatchers)
	if err != nil {
		return err
	}

	// 创建一个查询引擎
	qe := g.queryEngine(request.MaxResolutionSeconds)
	// 创建一个可查询的对象
	queryable := g.queryableCreate(
		request.EnableDedup,           // 是否启用去重
		request.ReplicaLabels,         // 复制标签
		storeMatchers,                 // 标签匹配器
		maxResolution,                 // 最大分辨率
		request.EnablePartialResponse, // 是否启用部分响应
		request.EnableQueryPushdown,   // 是否启用查询下推
		false,                         // 是否启用查询缓存
	)
	// 创建一个即时查询对象
	qry, err := qe.NewInstantQuery(queryable, request.Query, ts)
	if err != nil {
		return err
	}

	// 执行查询
	result := qry.Exec(ctx)
	// 发送查询警告
	if err := server.Send(querypb.NewQueryWarningsResponse(result.Warnings)); err != nil {
		return nil
	}

	// 根据查询结果的类型进行处理
	switch vector := result.Value.(type) {
	case promql.Scalar:
		// 处理标量类型的结果
		series := &prompb.TimeSeries{
			Samples: []prompb.Sample{{Value: vector.V, Timestamp: vector.T}},
		}
		if err := server.Send(querypb.NewQueryResponse(series)); err != nil {
			return err
		}
	case promql.Vector:
		// 处理向量类型的结果
		for _, sample := range vector {
			series := &prompb.TimeSeries{
				Labels:  labelpb.ZLabelsFromPromLabels(sample.Metric),
				Samples: prompb.SamplesFromPromqlPoints([]promql.Point{sample.Point}),
			}
			if err := server.Send(querypb.NewQueryResponse(series)); err != nil {
				return err
			}
		}

		return nil
	}

	return nil
}

func (g *GRPCAPI) QueryRange(request *querypb.QueryRangeRequest, srv querypb.Query_QueryRangeServer) error {
	// 创建一个背景上下文
	ctx := context.Background()
	// 如果请求中有超时时间设置，则创建带超时的上下文
	if request.TimeoutSeconds != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, time.Duration(request.TimeoutSeconds))
		defer cancel()
	}

	// 获取最大分辨率
	maxResolution := request.MaxResolutionSeconds
	// 如果没有设置最大分辨率，则使用默认值
	if request.MaxResolutionSeconds == 0 {
		maxResolution = g.defaultMaxResolutionSeconds.Milliseconds() / 1000
	}

	// 将请求中的StoreMatchers转换为LabelMatchers
	storeMatchers, err := querypb.StoreMatchersToLabelMatchers(request.StoreMatchers)
	if err != nil {
		return err
	}

	// 创建查询引擎
	qe := g.queryEngine(request.MaxResolutionSeconds)
	// 创建可查询对象
	queryable := g.queryableCreate(
		request.EnableDedup,           // 请求中的副本标签
		request.ReplicaLabels,         // 请求中的副本标签
		storeMatchers,                 // 请求中的StoreMatchers
		maxResolution,                 // 请求中的最大分辨率
		request.EnablePartialResponse, // 请求中的部分响应
		request.EnableQueryPushdown,   // 请求中的查询下推
		false,                         // 请求中的降采样
	)

	// 获取查询的起始和结束时间
	startTime := time.Unix(request.StartTimeSeconds, 0)
	endTime := time.Unix(request.EndTimeSeconds, 0)
	// 获取查询的时间间隔
	interval := time.Duration(request.IntervalSeconds) * time.Second

	// 创建范围查询
	qry, err := qe.NewRangeQuery(queryable, request.Query, startTime, endTime, interval)
	if err != nil {
		return err
	}

	// 执行查询
	result := qry.Exec(ctx)
	// 发送查询结果的警告信息
	if err := srv.Send(querypb.NewQueryRangeWarningsResponse(result.Warnings)); err != nil {
		return err
	}

	// 根据查询结果的类型进行处理
	switch matrix := result.Value.(type) {
	case promql.Matrix:
		// 遍历查询结果中的每个时间序列
		for _, series := range matrix {
			// 将时间序列转换为协议缓冲区格式
			series := &prompb.TimeSeries{
				Labels:  labelpb.ZLabelsFromPromLabels(series.Metric),  // 将Prometheus的标签转换为Label格式
				Samples: prompb.SamplesFromPromqlPoints(series.Points), // 将Prometheus的点转换为Sample格式
			}
			// 发送时间序列到客户端
			if err := srv.Send(querypb.NewQueryRangeResponse(series)); err != nil {
				return err
			}
		}

		return nil
	}

	return nil
}