package vm import ( "bufio" "compress/gzip" "context" "errors" "fmt" "io" "net/http" "strings" "sync" "time" "github.com/VictoriaMetrics/VictoriaMetrics/app/vmctl/backoff" "github.com/VictoriaMetrics/VictoriaMetrics/app/vmctl/barpool" "github.com/VictoriaMetrics/VictoriaMetrics/app/vmctl/limiter" "github.com/VictoriaMetrics/VictoriaMetrics/lib/decimal" "github.com/cheggaaa/pb/v3" ) // Config contains list of params to configure // the Importer type Config struct { // VictoriaMetrics address to perform import requests // --httpListenAddr value for single node version // --httpListenAddr value of vmselect component for cluster version Addr string // Transport allows specifying custom http.Transport Transport *http.Transport // Concurrency defines number of worker // performing the import requests concurrently Concurrency uint8 // Whether to apply gzip compression Compress bool // AccountID for cluster version. // Empty value assumes it is a single node version AccountID string // BatchSize defines how many samples // importer collects before sending the import request BatchSize int // User name for basic auth User string // Password for basic auth Password string // SignificantFigures defines the number of significant figures to leave // in metric values before importing. // Zero value saves all the significant decimal places SignificantFigures int // RoundDigits defines the number of decimal digits after the point that must be left // in metric values before importing. RoundDigits int // ExtraLabels that will be added to all imported series. Must be in label=value format. ExtraLabels []string // RateLimit defines a data transfer speed in bytes per second. // Is applied to each worker (see Concurrency) independently. RateLimit int64 // Whether to disable progress bar per VM worker DisableProgressBar bool } // Importer performs insertion of timeseries // via VictoriaMetrics import protocol // see https://docs.victoriametrics.com/#how-to-import-time-series-data type Importer struct { addr string client *http.Client importPath string compress bool user string password string close chan struct{} input chan *TimeSeries errors chan *ImportError rl *limiter.Limiter wg sync.WaitGroup once sync.Once s *stats backoff *backoff.Backoff } // ResetStats resets im stats. func (im *Importer) ResetStats() { im.s = &stats{ startTime: time.Now(), } } // Stats returns im stats. func (im *Importer) Stats() string { return im.s.String() } // AddExtraLabelsToImportPath - adds extra labels query params to given url path. func AddExtraLabelsToImportPath(path string, extraLabels []string) (string, error) { dst := path separator := "?" for _, extraLabel := range extraLabels { if !strings.Contains(extraLabel, "=") { return path, fmt.Errorf("bad format for extra_label flag, it must be `key=value`, got: %q", extraLabel) } if strings.Contains(dst, "?") { separator = "&" } dst += fmt.Sprintf("%sextra_label=%s", separator, extraLabel) } return dst, nil } // NewImporter creates new Importer for the given cfg. func NewImporter(ctx context.Context, cfg Config) (*Importer, error) { if cfg.Concurrency < 1 { return nil, fmt.Errorf("concurrency can't be lower than 1") } addr := strings.TrimRight(cfg.Addr, "/") // if single version // see https://docs.victoriametrics.com/#how-to-import-time-series-data importPath := addr + "/api/v1/import" if cfg.AccountID != "" { // if cluster version // see https://docs.victoriametrics.com/cluster-victoriametrics/#url-format importPath = fmt.Sprintf("%s/insert/%s/prometheus/api/v1/import", addr, cfg.AccountID) } importPath, err := AddExtraLabelsToImportPath(importPath, cfg.ExtraLabels) if err != nil { return nil, err } client := &http.Client{} if cfg.Transport != nil { client.Transport = cfg.Transport } im := &Importer{ addr: addr, client: client, importPath: importPath, compress: cfg.Compress, user: cfg.User, password: cfg.Password, rl: limiter.NewLimiter(cfg.RateLimit), close: make(chan struct{}), input: make(chan *TimeSeries, cfg.Concurrency*4), errors: make(chan *ImportError, cfg.Concurrency), backoff: backoff.New(), } if err := im.Ping(); err != nil { return nil, fmt.Errorf("ping to %q failed: %s", addr, err) } if cfg.BatchSize < 1 { cfg.BatchSize = 1e5 } im.wg.Add(int(cfg.Concurrency)) for i := 0; i < int(cfg.Concurrency); i++ { var bar *pb.ProgressBar if !cfg.DisableProgressBar { pbPrefix := fmt.Sprintf(`{{ green "VM worker %d:" }}`, i) bar = barpool.AddWithTemplate(pbPrefix+pbTpl, 0) } go func(bar *pb.ProgressBar) { defer im.wg.Done() im.startWorker(ctx, bar, cfg.BatchSize, cfg.SignificantFigures, cfg.RoundDigits) }(bar) } im.ResetStats() return im, nil } const pbTpl = `{{ (cycle . "←" "↖" "↑" "↗" "→" "↘" "↓" "↙" ) }} {{speed . "%s samples/s"}}` // ImportError is type of error generated // in case of unsuccessful import request type ImportError struct { // The batch of timeseries processed by importer at the moment Batch []*TimeSeries // The error that appeared during insert // If err is nil - no error happened and Batch // Is the latest delivered Batch. Err error } // Errors returns a channel for receiving // import errors if any func (im *Importer) Errors() chan *ImportError { return im.errors } // Input returns a channel for sending timeseries // that need to be imported func (im *Importer) Input(ts *TimeSeries) error { select { case <-im.close: return fmt.Errorf("importer is closed") case im.input <- ts: return nil case err := <-im.errors: if err != nil && err.Err != nil { return err.Err } return fmt.Errorf("process aborted") } } // Close sends signal to all goroutines to exit // and waits until they are finished func (im *Importer) Close() { im.once.Do(func() { close(im.close) close(im.input) im.wg.Wait() close(im.errors) }) } func (im *Importer) startWorker(ctx context.Context, bar *pb.ProgressBar, batchSize, significantFigures, roundDigits int) { var batch []*TimeSeries var dataPoints int var waitForBatch time.Time for { select { case <-im.close: for ts := range im.input { ts = roundTimeseriesValue(ts, significantFigures, roundDigits) batch = append(batch, ts) } exitErr := &ImportError{ Batch: batch, } retryableFunc := func() error { return im.Import(batch) } _, err := im.backoff.Retry(ctx, retryableFunc) if err != nil { exitErr.Err = err } im.errors <- exitErr return case ts, ok := <-im.input: if !ok { continue } // init waitForBatch when first // value was received if waitForBatch.IsZero() { waitForBatch = time.Now() } ts = roundTimeseriesValue(ts, significantFigures, roundDigits) batch = append(batch, ts) dataPoints += len(ts.Values) if bar != nil { bar.Add(len(ts.Values)) } if dataPoints < batchSize { continue } im.s.Lock() im.s.idleDuration += time.Since(waitForBatch) im.s.Unlock() if err := im.flush(ctx, batch); err != nil { im.errors <- &ImportError{ Batch: batch, Err: err, } // make a new batch, since old one was referenced as err batch = make([]*TimeSeries, len(batch)) } dataPoints = 0 batch = batch[:0] waitForBatch = time.Now() } } } func (im *Importer) flush(ctx context.Context, b []*TimeSeries) error { retryableFunc := func() error { return im.Import(b) } attempts, err := im.backoff.Retry(ctx, retryableFunc) if err != nil { return fmt.Errorf("import failed with %d retries: %s", attempts, err) } im.s.Lock() im.s.retries = attempts im.s.Unlock() return nil } // Ping sends a ping to im.addr. func (im *Importer) Ping() error { url := fmt.Sprintf("%s/health", im.addr) req, err := http.NewRequest(http.MethodGet, url, nil) if err != nil { return fmt.Errorf("cannot create request to %q: %s", im.addr, err) } if im.user != "" { req.SetBasicAuth(im.user, im.password) } resp, err := im.client.Do(req) if err != nil { return err } if resp.StatusCode != http.StatusOK { return fmt.Errorf("bad status code: %d", resp.StatusCode) } return nil } // Import imports tsBatch. func (im *Importer) Import(tsBatch []*TimeSeries) error { if len(tsBatch) < 1 { return nil } pr, pw := io.Pipe() req, err := http.NewRequest(http.MethodPost, im.importPath, pr) if err != nil { return fmt.Errorf("cannot create request to %q: %s", im.addr, err) } if im.user != "" { req.SetBasicAuth(im.user, im.password) } if im.compress { req.Header.Set("Content-Encoding", "gzip") } errCh := make(chan error) go func() { errCh <- im.do(req) close(errCh) }() w := io.Writer(pw) if im.compress { zw, err := gzip.NewWriterLevel(w, 1) if err != nil { return fmt.Errorf("unexpected error when creating gzip writer: %s", err) } w = zw } w = limiter.NewWriteLimiter(w, im.rl) bw := bufio.NewWriterSize(w, 16*1024) var totalSamples, totalBytes int for _, ts := range tsBatch { n, err := ts.write(bw) if err != nil { return fmt.Errorf("write err: %w", err) } totalBytes += n totalSamples += len(ts.Values) } if err := bw.Flush(); err != nil { return err } if closer, ok := w.(io.Closer); ok { err := closer.Close() if err != nil { return err } } if err := pw.Close(); err != nil { return err } requestErr := <-errCh if requestErr != nil { return fmt.Errorf("import request error for %q: %w", im.addr, requestErr) } im.s.Lock() im.s.bytes += uint64(totalBytes) im.s.samples += uint64(totalSamples) im.s.requests++ im.s.Unlock() return nil } // ErrBadRequest represents bad request error. var ErrBadRequest = errors.New("bad request") func (im *Importer) do(req *http.Request) error { resp, err := im.client.Do(req) if err != nil { return fmt.Errorf("unexpected error when performing request: %s", err) } defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusNoContent { body, err := io.ReadAll(resp.Body) if err != nil { return fmt.Errorf("failed to read response body for status code %d: %s", resp.StatusCode, err) } if resp.StatusCode == http.StatusBadRequest { return fmt.Errorf("%w: unexpected response code %d: %s", ErrBadRequest, resp.StatusCode, string(body)) } return fmt.Errorf("unexpected response code %d: %s", resp.StatusCode, string(body)) } return nil } func byteCountSI(b int64) string { const unit = 1000 if b < unit { return fmt.Sprintf("%d B", b) } div, exp := int64(unit), 0 for n := b / unit; n >= unit; n /= unit { div *= unit exp++ } return fmt.Sprintf("%.1f %cB", float64(b)/float64(div), "kMGTPE"[exp]) } func roundTimeseriesValue(ts *TimeSeries, significantFigures, roundDigits int) *TimeSeries { if significantFigures > 0 { for i, v := range ts.Values { ts.Values[i] = decimal.RoundToSignificantFigures(v, significantFigures) } } if roundDigits < 100 { for i, v := range ts.Values { ts.Values[i] = decimal.RoundToDecimalDigits(v, roundDigits) } } return ts }