gotidb/pkg/engine/bolt/bolt.go

package bolt

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"git.pyer.club/kingecg/gotidb/pkg/engine"
	bolt "go.etcd.io/bbolt"
)

const (
	// bucket 名称
	seriesBucketPrefix = "series_" // 序列数据 bucket 前缀
	metaBucketName     = "meta"    // 元数据 bucket
	indexBucketName    = "index"   // 索引 bucket

	// 元数据键
	statsKey = "stats" // 统计信息键
)

// BoltEngine 实现基于 BoltDB 的存储引擎
type BoltEngine struct {
	mu     sync.RWMutex
	config *engine.BoltEngineConfig
	db     *bolt.DB
	stats  engine.EngineStats
	opened bool
	closed bool
}

// 确保 BoltEngine 实现了 Engine 接口
var _ engine.Engine = (*BoltEngine)(nil)

// NewBoltEngine 创建一个新的 Bolt 引擎
func NewBoltEngine(config *engine.BoltEngineConfig) (engine.Engine, error) {
	if config == nil {
		return nil, fmt.Errorf("bolt engine config cannot be nil")
	}

	// 确保数据目录存在
	if err := os.MkdirAll(filepath.Dir(config.Path), 0755); err != nil {
		return nil, fmt.Errorf("failed to create data directory: %v", err)
	}

	return &BoltEngine{
		config: config,
		stats: engine.EngineStats{
			LastWriteTime: time.Now(),
		},
	}, nil
}

// Open 实现 Engine 接口
func (b *BoltEngine) Open() error {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.opened {
		return fmt.Errorf("bolt engine already opened")
	}

	if b.closed {
		return fmt.Errorf("bolt engine already closed")
	}

	// 打开数据库
	db, err := bolt.Open(b.config.Path, 0600, &bolt.Options{
		Timeout: 1 * time.Second,
	})
	if err != nil {
		return fmt.Errorf("failed to open bolt database: %v", err)
	}

	// 创建必要的 bucket
	err = db.Update(func(tx *bolt.Tx) error {
		// 创建元数据 bucket
		if _, err := tx.CreateBucketIfNotExists([]byte(metaBucketName)); err != nil {
			return fmt.Errorf("failed to create meta bucket: %v", err)
		}

		// 创建索引 bucket
		if _, err := tx.CreateBucketIfNotExists([]byte(indexBucketName)); err != nil {
			return fmt.Errorf("failed to create index bucket: %v", err)
		}

		return nil
	})
	if err != nil {
		db.Close()
		return fmt.Errorf("failed to create buckets: %v", err)
	}

	b.db = db
	b.opened = true

	// 加载统计信息
	if err := b.loadStats(); err != nil {
		return fmt.Errorf("failed to load stats: %v", err)
	}

	return nil
}

// Close 实现 Engine 接口
func (b *BoltEngine) Close() error {
	b.mu.Lock()
	defer b.mu.Unlock()

	if !b.opened {
		return fmt.Errorf("bolt engine not opened")
	}

	if b.closed {
		return fmt.Errorf("bolt engine already closed")
	}

	// 保存统计信息
	if err := b.saveStats(); err != nil {
		return fmt.Errorf("failed to save stats: %v", err)
	}

	// 关闭数据库
	if err := b.db.Close(); err != nil {
		return fmt.Errorf("failed to close bolt database: %v", err)
	}

	b.closed = true
	return nil
}

// Write 实现 Engine 接口
func (b *BoltEngine) Write(ctx context.Context, points []engine.DataPoint) error {
	if len(points) == 0 {
		return nil
	}

	b.mu.Lock()
	defer b.mu.Unlock()

	if !b.opened || b.closed {
		return fmt.Errorf("bolt engine not open")
	}

	startTime := time.Now()

	// 写入数据点
	err := b.db.Update(func(tx *bolt.Tx) error {
		for _, point := range points {
			seriesID := point.GetSeriesID()
			if seriesID == "" {
				b.stats.WriteErrors++
				continue
			}

			// 获取或创建序列 bucket
			bucketName := seriesBucketPrefix + seriesID
			bucket, err := tx.CreateBucketIfNotExists([]byte(bucketName))
			if err != nil {
				return fmt.Errorf("failed to create series bucket: %v", err)
			}

			// 序列化数据点
			value, err := json.Marshal(point)
			if err != nil {
				return fmt.Errorf("failed to marshal data point: %v", err)
			}

			// 使用时间戳作为键
			key := make([]byte, 8)
			binary.BigEndian.PutUint64(key, uint64(point.Timestamp))

			// 写入数据点
			if err := bucket.Put(key, value); err != nil {
				return fmt.Errorf("failed to write data point: %v", err)
			}

			// 更新索引
			indexBucket := tx.Bucket([]byte(indexBucketName))
			if err := indexBucket.Put([]byte(seriesID), []byte(seriesID)); err != nil {
				return fmt.Errorf("failed to update index: %v", err)
			}
		}

		return nil
	})

	if err != nil {
		return fmt.Errorf("failed to write points: %v", err)
	}

	// 更新统计信息
	b.stats.PointsCount += int64(len(points))
	b.stats.LastWriteTime = time.Now()
	b.stats.WriteLatency = time.Since(startTime) / time.Duration(len(points))

	// 更新序列数
	if err := b.updateSeriesCount(); err != nil {
		return fmt.Errorf("failed to update series count: %v", err)
	}

	return nil
}

// Stats 实现 Engine 接口
func (b *BoltEngine) Stats() engine.EngineStats {
	b.mu.RLock()
	defer b.mu.RUnlock()

	stats := b.stats

	// 获取数据库大小
	if b.db != nil {
		stats.DiskUsage = b.db.Stats().TxStats.PageCount * int64(b.db.Info().PageSize)
	}

	return stats
}

// 注册 Bolt 引擎创建函数
func init() {
	engine.RegisterBoltEngine(NewBoltEngine)
}