sig-pub/pkg/storage/tsdb/victoria_metrics/vm.go

package vmts

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"runtime/debug"
	"sig-pub/pkg/config"
	"sig-pub/pkg/types"
	"sig-pub/pkg/utils/collect"
	"sig-pub/pkg/zlog"

	"github.com/bytedance/sonic"
	"github.com/govalues/decimal"
	"github.com/klauspost/compress/zstd"
)

// VictoriaMetricsTSDB 时序库
type VictoriaMetricsTSDB struct {
	addr string
}

func NewVictoriaMetricsTSDB(conf config.VictoriaMetricsConfig) *VictoriaMetricsTSDB {
	return &VictoriaMetricsTSDB{
		addr: conf.Addr,
	}
}

func (vm *VictoriaMetricsTSDB) SaveKlines(inst types.TradeInstance, klines []*types.Kline) (err error) {
	metrics := Kline2Metrics(inst, klines)
	err = vm.batchWriteMetrics(metrics)
	return
}

func (vm *VictoriaMetricsTSDB) batchWriteMetrics(metrics []*Metric) (err error) {
	var buf bytes.Buffer
	// gz := gzip.NewWriter(&buf)
	var data []byte
	for _, metric := range metrics {
		data, err = metric.ToRowJson()
		if err != nil {
			return
		}
		buf.Write(data)
		buf.Write([]byte("\n"))
		// gz.Write(data)
		// gz.Write([]byte("\n"))
	}
	// gz.Close()

	// file, _ := os.Open(fmt.Sprintf("%d.json", time.Now().Unix()))
	// defer file.Close()
	// err = os.WriteFile(fmt.Sprintf("./%d.json", time.Now().Unix()), buf.Bytes(), os.ModeAppend)
	// if err != nil {
	// 	return
	// }

	// datas, err := compressData(buf.Bytes())
	// {kind="high"}[30m]

	resp, err := http.Post(fmt.Sprintf("%s/api/v1/import", vm.addr), "application/json", &buf)
	if err != nil {
		return
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		zlog.Info("batch import status error: ", resp.Status)
	}
	return
}

func compressData(data []byte) ([]byte, error) {
	var b bytes.Buffer
	encoder, _ := zstd.NewWriter(&b)
	defer encoder.Close()

	if _, err := encoder.Write(data); err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

// 获取原始k线列表 [end, ..., start]
func (vm *VictoriaMetricsTSDB) ListRangeKline(inst types.TradeInstance, interval types.Interval, start, end int64) (klines []*types.Kline, err error) {
	defer func() {
		if r := recover(); r != nil {
			zlog.Error("vmdb get range kline recover error:", r)
			debug.PrintStack()
			err = fmt.Errorf("%v", r)
		}
	}()
	if inst.InstId == "" {
		err = errors.New("instid is empty")
		return
	}
	if start <= 0 || end <= 0 {
		err = errors.New("invalid time range")
		return
	}

	match := fmt.Sprintf("%s{exchange=\"%s\", interval=\"%s\"}", inst.InstId, inst.Exchange.String(), interval)
	params := fmt.Sprintf("start=%d&end=%d&match[]=%s", start, end, url.QueryEscape(match))
	resp, err := http.Get(fmt.Sprintf("%s/api/v1/export?%s", vm.addr, params))
	if err != nil {
		return
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			zlog.Error("close vmtsdb response error:", err)
		}
	}()

	// read response json line
	reader := bufio.NewReader(resp.Body)
	var vmLineBytes []byte
	for {
		// todo 优化
		vmLineBytes, err = reader.ReadBytes('\n')
		if err == io.EOF || len(vmLineBytes) == 0 {
			err = nil
			break
		}
		if err != nil {
			zlog.Error(err)
			return
		}

		vmMetric := new(VMMetricKline)
		if err = sonic.Unmarshal(vmLineBytes, vmMetric); err != nil {
			return
		}
		for i, ts := range vmMetric.Timestamps {
			if len(klines) <= i {
				klines = append(klines, &types.Kline{
					Interval: types.Interval(vmMetric.Metric.Interval),
					Ts:       ts,
					Confirm:  true,
				})
			}
			kline := klines[i]
			if kline.Ts != ts {
				err = errors.New("vm metric kline integrate error")
				return
			}
			value := vmMetric.Values[i]
			switch vmMetric.Metric.Kind {
			case "open":
				kline.Open = value
			case "close":
				kline.Close = value
			case "high":
				kline.High = value
			case "low":
				kline.Low = value
			case "vol":
				kline.Vol = value
			case "volQuote":
				kline.VolQuote = value
			}
		}
		// zlog.Infof("response body line: %#v", vmMetric)
	}
	collect.Reverse(klines)
	// for _, kline := range klines {
	// 	zlog.Infof("kline: %#v", kline)
	// }
	return
}

type VMMetricKline struct {
	Metric struct {
		Name     string `json:"__name__"`
		Interval string `json:"interval"`
		Kind     string `json:"kind"`
	} `json:"metric"`
	Values     []decimal.Decimal `json:"values"`
	Timestamps []int64           `json:"timestamps"`
}