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

package vmts

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

	"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线列表
func (vm *VictoriaMetricsTSDB) GetRangeKline(inst types.TradeInstance, interval types.Interval, start, end int64) (err error) {
	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{interval=\"%s\"}", inst.InstId, 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)
	for {
		line, err2 := reader.ReadBytes('\n')
		if err2 == io.EOF {
			break
		}
		if err2 != nil {
			zlog.Error(err)
			err = err2
			return
		}
		zlog.Infof("response body line: %s", string(line))
	}
	return
}