influxdb/cmd/influxd/launcher/engine.go

package launcher

import (
	"context"
	"io"
	"os"
	"sync"
	"time"

	"github.com/influxdata/influxql"

	"github.com/influxdata/influxdb/v2"
	"github.com/influxdata/influxdb/v2/http"
	"github.com/influxdata/influxdb/v2/kit/platform"
	"github.com/influxdata/influxdb/v2/kit/prom"
	"github.com/influxdata/influxdb/v2/models"
	"github.com/influxdata/influxdb/v2/storage"
	"github.com/prometheus/client_golang/prometheus"
	"go.uber.org/zap"
)

var _ Engine = (*storage.Engine)(nil)

// Engine defines the time-series storage engine.  Wraps *storage.Engine
// to facilitate testing.
type Engine interface {
	influxdb.DeleteService
	storage.PointsWriter
	storage.EngineSchema
	prom.PrometheusCollector
	influxdb.BackupService
	influxdb.RestoreService

	SeriesCardinality(ctx context.Context, bucketID platform.ID) int64

	TSDBStore() storage.TSDBStore
	MetaClient() storage.MetaClient

	WithLogger(log *zap.Logger)
	Open(context.Context) error
	Close() error
}

var _ Engine = (*TemporaryEngine)(nil)
var _ http.Flusher = (*TemporaryEngine)(nil)

// TemporaryEngine creates a time-series storage engine backed
// by a temporary directory that is removed on Close.
type TemporaryEngine struct {
	path    string
	config  storage.Config
	options []storage.Option

	mu     sync.Mutex
	opened bool

	engine    *storage.Engine
	tsdbStore temporaryTSDBStore

	log *zap.Logger
}

// NewTemporaryEngine creates a new engine that places the storage engine files into
// a temporary directory; used for testing.
func NewTemporaryEngine(c storage.Config, options ...storage.Option) *TemporaryEngine {
	return &TemporaryEngine{
		config:  c,
		options: options,
		log:     zap.NewNop(),
	}
}

// Open creates a temporary directory and opens the engine.
func (t *TemporaryEngine) Open(ctx context.Context) error {
	t.mu.Lock()
	defer t.mu.Unlock()

	if t.opened {
		return nil
	}

	path, err := os.MkdirTemp("", "e2e")
	if err != nil {
		return err
	}

	t.path = path
	t.engine = storage.NewEngine(path, t.config, t.options...)
	t.engine.WithLogger(t.log)

	if err := t.engine.Open(ctx); err != nil {
		_ = os.RemoveAll(path)
		return err
	}

	t.tsdbStore.TSDBStore = t.engine.TSDBStore()

	t.opened = true
	return nil
}

// Close will remove the directory containing the time-series files.
func (t *TemporaryEngine) Close() error {
	t.mu.Lock()
	defer t.mu.Unlock()

	t.opened = false
	err := t.engine.Close()
	_ = os.RemoveAll(t.path)
	return err
}

// WritePoints stores points into the storage engine.
func (t *TemporaryEngine) WritePoints(ctx context.Context, orgID platform.ID, bucketID platform.ID, points []models.Point) error {
	return t.engine.WritePoints(ctx, orgID, bucketID, points)
}

// SeriesCardinality returns the number of series in the engine.
func (t *TemporaryEngine) SeriesCardinality(ctx context.Context, bucketID platform.ID) int64 {
	return t.engine.SeriesCardinality(ctx, bucketID)
}

// DeleteBucketRangePredicate will delete a bucket from the range and predicate.
func (t *TemporaryEngine) DeleteBucketRangePredicate(ctx context.Context, orgID, bucketID platform.ID, min, max int64, pred influxdb.Predicate, measurement influxql.Expr) error {
	return t.engine.DeleteBucketRangePredicate(ctx, orgID, bucketID, min, max, pred, measurement)
}

func (t *TemporaryEngine) CreateBucket(ctx context.Context, b *influxdb.Bucket) error {
	return t.engine.CreateBucket(ctx, b)
}

func (t *TemporaryEngine) UpdateBucketRetentionPolicy(ctx context.Context, bucketID platform.ID, upd *influxdb.BucketUpdate) error {
	return t.engine.UpdateBucketRetentionPolicy(ctx, bucketID, upd)
}

// DeleteBucket deletes a bucket from the time-series data.
func (t *TemporaryEngine) DeleteBucket(ctx context.Context, orgID, bucketID platform.ID) error {
	return t.engine.DeleteBucket(ctx, orgID, bucketID)
}

// WithLogger sets the logger on the engine. It must be called before Open.
func (t *TemporaryEngine) WithLogger(log *zap.Logger) {
	t.log = log.With(zap.String("service", "temporary_engine"))
}

// PrometheusCollectors returns all the prometheus collectors associated with
// the engine and its components.
func (t *TemporaryEngine) PrometheusCollectors() []prometheus.Collector {
	return t.engine.PrometheusCollectors()
}

// Flush will remove the time-series files and re-open the engine.
func (t *TemporaryEngine) Flush(ctx context.Context) {
	if err := t.Close(); err != nil {
		t.log.Fatal("unable to close engine", zap.Error(err))
	}

	if err := t.Open(ctx); err != nil {
		t.log.Fatal("unable to open engine", zap.Error(err))
	}
}

func (t *TemporaryEngine) BackupKVStore(ctx context.Context, w io.Writer) error {
	return t.engine.BackupKVStore(ctx, w)
}

func (t *TemporaryEngine) RLockKVStore() {
	t.engine.RLockKVStore()
}

func (t *TemporaryEngine) RUnlockKVStore() {
	t.engine.RUnlockKVStore()
}

func (t *TemporaryEngine) RestoreKVStore(ctx context.Context, r io.Reader) error {
	return t.engine.RestoreKVStore(ctx, r)
}

func (t *TemporaryEngine) RestoreBucket(ctx context.Context, id platform.ID, dbi []byte) (map[uint64]uint64, error) {
	return t.engine.RestoreBucket(ctx, id, dbi)
}

func (t *TemporaryEngine) BackupShard(ctx context.Context, w io.Writer, shardID uint64, since time.Time) error {
	return t.engine.BackupShard(ctx, w, shardID, since)
}

func (t *TemporaryEngine) RestoreShard(ctx context.Context, shardID uint64, r io.Reader) error {
	return t.engine.RestoreShard(ctx, shardID, r)
}

func (t *TemporaryEngine) TSDBStore() storage.TSDBStore {
	return &t.tsdbStore
}

func (t *TemporaryEngine) MetaClient() storage.MetaClient {
	return t.engine.MetaClient()
}

type temporaryTSDBStore struct {
	storage.TSDBStore
}