341 lines
8.5 KiB
Go
341 lines
8.5 KiB
Go
package vmclient
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"net/url"
|
||
"time"
|
||
)
|
||
|
||
// VMClient VictoriaMetrics 客户端
|
||
type VMClient struct {
|
||
baseURL string
|
||
httpClient *http.Client
|
||
writeTimeout time.Duration
|
||
queryTimeout time.Duration
|
||
}
|
||
|
||
// QueryResult 查询结果
|
||
type QueryResult struct {
|
||
Status string `json:"status"`
|
||
Data ResultData `json:"data"`
|
||
}
|
||
|
||
// ResultData 查询结果数据
|
||
type ResultData struct {
|
||
ResultType string `json:"resultType"`
|
||
Result []Result `json:"result"`
|
||
}
|
||
|
||
// Result 单个时间序列结果
|
||
type Result struct {
|
||
Metric map[string]string `json:"metric"`
|
||
Values [][]interface{} `json:"values"` // [[timestamp, value], ...]
|
||
}
|
||
|
||
// DataPoint 数据点
|
||
type DataPoint struct {
|
||
Timestamp int64
|
||
Value float64
|
||
Labels map[string]string
|
||
}
|
||
|
||
// Metric VictoriaMetrics JSON Line Format 指标
|
||
type Metric struct {
|
||
Metric map[string]string `json:"metric"`
|
||
Values []float64 `json:"values"`
|
||
Timestamps []int64 `json:"timestamps"`
|
||
}
|
||
|
||
// NewVMClient 创建 VictoriaMetrics 客户端
|
||
func NewVMClient(baseURL string, writeTimeout, queryTimeout time.Duration) *VMClient {
|
||
if writeTimeout == 0 {
|
||
writeTimeout = 30 * time.Second
|
||
}
|
||
if queryTimeout == 0 {
|
||
queryTimeout = 60 * time.Second
|
||
}
|
||
|
||
return &VMClient{
|
||
baseURL: baseURL,
|
||
httpClient: &http.Client{
|
||
Timeout: writeTimeout, // 默认超时
|
||
},
|
||
writeTimeout: writeTimeout,
|
||
queryTimeout: queryTimeout,
|
||
}
|
||
}
|
||
|
||
// Write 写入指标（VictoriaMetrics JSON Line Format）
|
||
func (c *VMClient) Write(ctx context.Context, metrics []Metric) error {
|
||
if len(metrics) == 0 {
|
||
return nil
|
||
}
|
||
|
||
reqCtx, cancel := context.WithTimeout(ctx, c.writeTimeout)
|
||
defer cancel()
|
||
|
||
// 将 Metric 数组转换为 JSON Line Format (NDJSON)
|
||
var buf bytes.Buffer
|
||
encoder := json.NewEncoder(&buf)
|
||
for _, metric := range metrics {
|
||
if err := encoder.Encode(metric); err != nil {
|
||
return fmt.Errorf("encode metric failed: %w", err)
|
||
}
|
||
}
|
||
|
||
req, err := http.NewRequestWithContext(reqCtx, "POST", c.baseURL+"/api/v1/import", &buf)
|
||
if err != nil {
|
||
return fmt.Errorf("create request failed: %w", err)
|
||
}
|
||
|
||
req.Header.Set("Content-Type", "application/x-ndjson")
|
||
|
||
resp, err := c.httpClient.Do(req)
|
||
if err != nil {
|
||
return fmt.Errorf("write metrics failed: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusNoContent && resp.StatusCode != http.StatusOK {
|
||
body, _ := io.ReadAll(resp.Body)
|
||
return fmt.Errorf("write metrics failed with status %d: %s", resp.StatusCode, string(body))
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// AutoStep 根据查询范围自动生成 step（适用于 VictoriaMetrics）
|
||
func AutoStep(start, end time.Time) time.Duration {
|
||
r := end.Sub(start)
|
||
|
||
switch {
|
||
case r <= time.Hour:
|
||
return 10 * time.Second
|
||
case r <= 3*time.Hour:
|
||
return 15 * time.Second
|
||
case r <= 6*time.Hour:
|
||
return 30 * time.Second
|
||
case r <= 12*time.Hour:
|
||
return time.Minute
|
||
case r <= 24*time.Hour:
|
||
return 2 * time.Minute
|
||
case r <= 3*24*time.Hour:
|
||
return 5 * time.Minute
|
||
case r <= 7*24*time.Hour:
|
||
return 10 * time.Minute
|
||
case r <= 30*24*time.Hour:
|
||
return 30 * time.Minute
|
||
default:
|
||
return time.Hour
|
||
}
|
||
}
|
||
|
||
// QueryRange 范围查询
|
||
// 如果 step 为 0，则让 VictoriaMetrics 自动选择合适的步长
|
||
func (c *VMClient) QueryRange(ctx context.Context, query string, start, end time.Time, step time.Duration) (*QueryResult, error) {
|
||
reqCtx, cancel := context.WithTimeout(ctx, c.queryTimeout)
|
||
defer cancel()
|
||
|
||
params := url.Values{}
|
||
params.Set("query", query)
|
||
params.Set("start", fmt.Sprintf("%d", start.Unix()))
|
||
params.Set("end", fmt.Sprintf("%d", end.Unix()))
|
||
if step > 0 {
|
||
params.Set("step", fmt.Sprintf("%ds", int(step.Seconds())))
|
||
} else {
|
||
autoStep := AutoStep(start, end)
|
||
params.Set("step", fmt.Sprintf("%ds", int(autoStep.Seconds())))
|
||
}
|
||
|
||
reqURL := fmt.Sprintf("%s/api/v1/query_range?%s", c.baseURL, params.Encode())
|
||
|
||
req, err := http.NewRequestWithContext(reqCtx, "GET", reqURL, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create request failed: %w", err)
|
||
}
|
||
|
||
resp, err := c.httpClient.Do(req)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("query range failed: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
body, _ := io.ReadAll(resp.Body)
|
||
return nil, fmt.Errorf("query range failed with status %d: %s", resp.StatusCode, string(body))
|
||
}
|
||
|
||
var result QueryResult
|
||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||
return nil, fmt.Errorf("decode response failed: %w", err)
|
||
}
|
||
|
||
if result.Status != "success" {
|
||
return nil, fmt.Errorf("query failed with status: %s", result.Status)
|
||
}
|
||
|
||
return &result, nil
|
||
}
|
||
|
||
// Query 即时查询
|
||
func (c *VMClient) Query(ctx context.Context, query string) (*QueryResult, error) {
|
||
reqCtx, cancel := context.WithTimeout(ctx, c.queryTimeout)
|
||
defer cancel()
|
||
|
||
params := url.Values{}
|
||
params.Set("query", query)
|
||
|
||
reqURL := fmt.Sprintf("%s/api/v1/query?%s", c.baseURL, params.Encode())
|
||
|
||
req, err := http.NewRequestWithContext(reqCtx, "GET", reqURL, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create request failed: %w", err)
|
||
}
|
||
|
||
resp, err := c.httpClient.Do(req)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("query failed: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
body, _ := io.ReadAll(resp.Body)
|
||
return nil, fmt.Errorf("query failed with status %d: %s", resp.StatusCode, string(body))
|
||
}
|
||
|
||
var result QueryResult
|
||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||
return nil, fmt.Errorf("decode response failed: %w", err)
|
||
}
|
||
|
||
if result.Status != "success" {
|
||
return nil, fmt.Errorf("query failed with status: %s", result.Status)
|
||
}
|
||
|
||
return &result, nil
|
||
}
|
||
|
||
// DeleteSeries 删除时间序列数据
|
||
func (c *VMClient) DeleteSeries(ctx context.Context, matchers []string) error {
|
||
if len(matchers) == 0 {
|
||
return nil
|
||
}
|
||
|
||
reqCtx, cancel := context.WithTimeout(ctx, c.queryTimeout)
|
||
defer cancel()
|
||
|
||
params := url.Values{}
|
||
for _, matcher := range matchers {
|
||
params.Add("match[]", matcher)
|
||
}
|
||
|
||
reqURL := fmt.Sprintf("%s/api/v1/admin/tsdb/delete_series?%s", c.baseURL, params.Encode())
|
||
req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, reqURL, nil)
|
||
if err != nil {
|
||
return fmt.Errorf("create request failed: %w", err)
|
||
}
|
||
|
||
resp, err := c.httpClient.Do(req)
|
||
if err != nil {
|
||
return fmt.Errorf("delete series failed: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNoContent {
|
||
body, _ := io.ReadAll(resp.Body)
|
||
return fmt.Errorf("delete series failed with status %d: %s", resp.StatusCode, string(body))
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// ConvertToDataPoints 将查询结果转换为数据点列表
|
||
func ConvertToDataPoints(result *QueryResult) []DataPoint {
|
||
if result == nil || len(result.Data.Result) == 0 {
|
||
return []DataPoint{}
|
||
}
|
||
|
||
var points []DataPoint
|
||
for _, r := range result.Data.Result {
|
||
for _, v := range r.Values {
|
||
if len(v) < 2 {
|
||
continue
|
||
}
|
||
|
||
// timestamp 是 float64（Unix 秒）
|
||
timestamp, ok := v[0].(float64)
|
||
if !ok {
|
||
continue
|
||
}
|
||
|
||
// value 是 string
|
||
valueStr, ok := v[1].(string)
|
||
if !ok {
|
||
continue
|
||
}
|
||
|
||
var value float64
|
||
if _, err := fmt.Sscanf(valueStr, "%f", &value); err != nil {
|
||
continue
|
||
}
|
||
|
||
points = append(points, DataPoint{
|
||
Timestamp: int64(timestamp * 1000), // 转换为毫秒
|
||
Value: value,
|
||
Labels: r.Metric,
|
||
})
|
||
}
|
||
}
|
||
|
||
return points
|
||
}
|
||
|
||
// GetLabelValues 获取指定 label 的所有值
|
||
func (c *VMClient) GetLabelValues(ctx context.Context, labelName string, match []string) ([]string, error) {
|
||
reqCtx, cancel := context.WithTimeout(ctx, c.queryTimeout)
|
||
defer cancel()
|
||
|
||
params := url.Values{}
|
||
for _, m := range match {
|
||
params.Add("match[]", m)
|
||
}
|
||
|
||
reqURL := fmt.Sprintf("%s/api/v1/label/%s/values?%s", c.baseURL, labelName, params.Encode())
|
||
|
||
req, err := http.NewRequestWithContext(reqCtx, "GET", reqURL, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create request failed: %w", err)
|
||
}
|
||
|
||
resp, err := c.httpClient.Do(req)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("get label values failed: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
body, _ := io.ReadAll(resp.Body)
|
||
return nil, fmt.Errorf("get label values failed with status %d: %s", resp.StatusCode, string(body))
|
||
}
|
||
|
||
var result struct {
|
||
Status string `json:"status"`
|
||
Data []string `json:"data"`
|
||
}
|
||
|
||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||
return nil, fmt.Errorf("decode response failed: %w", err)
|
||
}
|
||
|
||
if result.Status != "success" {
|
||
return nil, fmt.Errorf("get label values failed with status: %s", result.Status)
|
||
}
|
||
|
||
return result.Data, nil
|
||
}
|