gotidb/pkg/engine/utils.go

package engine

import (
	"context"
	"hash/fnv"
	"sync"
)

// WriteBuffer 实现写入缓冲区，用于合并小批量写入
type WriteBuffer struct {
	points  map[string][]DataPoint // 按序列ID分组的数据点
	mu      sync.Mutex
	maxSize int           // 缓冲区最大大小
	flushCh chan struct{} // 触发刷新的通道
	engine  Engine        // 底层存储引擎
}

// NewWriteBuffer 创建新的写入缓冲区
func NewWriteBuffer(engine Engine, maxSize int) *WriteBuffer {
	wb := &WriteBuffer{
		points:  make(map[string][]DataPoint),
		maxSize: maxSize,
		flushCh: make(chan struct{}, 1),
		engine:  engine,
	}
	return wb
}

// Add 添加数据点到缓冲区
func (wb *WriteBuffer) Add(point DataPoint) error {
	wb.mu.Lock()
	seriesID := point.SeriesID()
	wb.points[seriesID] = append(wb.points[seriesID], point)
	size := len(wb.points)
	wb.mu.Unlock()

	if size >= wb.maxSize {
		return wb.Flush()
	}
	return nil
}

// Flush 将缓冲区数据写入底层引擎
func (wb *WriteBuffer) Flush() error {
	wb.mu.Lock()
	points := make([]DataPoint, 0, len(wb.points)*10) // 预估大小
	for _, seriesPoints := range wb.points {
		points = append(points, seriesPoints...)
	}
	wb.points = make(map[string][]DataPoint)
	wb.mu.Unlock()

	if len(points) > 0 {
		return wb.engine.WriteBatch(context.Background(), points)
	}
	return nil
}

// Close 关闭写入缓冲区
func (wb *WriteBuffer) Close() error {
	return wb.Flush()
}

// ShardedLock 实现分片锁，用于减少锁竞争
type ShardedLock struct {
	locks     []sync.RWMutex
	shardMask uint64
}

// NewShardedLock 创建新的分片锁
func NewShardedLock(shards int) *ShardedLock {
	// 确保分片数是2的幂
	shards = nextPowerOfTwo(shards)
	return &ShardedLock{
		locks:     make([]sync.RWMutex, shards),
		shardMask: uint64(shards - 1),
	}
}

// getLockForKey 获取指定键的锁
func (sl *ShardedLock) getLockForKey(key string) *sync.RWMutex {
	h := fnv.New64()
	h.Write([]byte(key))
	hashVal := h.Sum64()
	return &sl.locks[hashVal&sl.shardMask]
}

// Lock 对指定键加写锁
func (sl *ShardedLock) Lock(key string) {
	sl.getLockForKey(key).Lock()
}

// Unlock 对指定键解除写锁
func (sl *ShardedLock) Unlock(key string) {
	sl.getLockForKey(key).Unlock()
}

// RLock 对指定键加读锁
func (sl *ShardedLock) RLock(key string) {
	sl.getLockForKey(key).RLock()
}

// RUnlock 对指定键解除读锁
func (sl *ShardedLock) RUnlock(key string) {
	sl.getLockForKey(key).RUnlock()
}

// CompactTimeSeriesBlock 实现时序数据的紧凑存储
type CompactTimeSeriesBlock struct {
	baseTime    int64  // 基准时间戳
	deltaEncode []byte // 使用delta编码存储时间戳
	values      []byte // 压缩存储的值
}

// NewCompactBlock 创建新的紧凑存储块
func NewCompactBlock(baseTime int64, capacity int) *CompactTimeSeriesBlock {
	return &CompactTimeSeriesBlock{
		baseTime:    baseTime,
		deltaEncode: make([]byte, 0, capacity*8), // 预留足够空间
		values:      make([]byte, 0, capacity*8),
	}
}

// nextPowerOfTwo 返回大于等于n的最小2的幂
func nextPowerOfTwo(n int) int {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n++
	return n
}

// TimeRangeIndex 实现时间范围索引
type TimeRangeIndex struct {
	windows   []timeWindow
	blockSize int64 // 时间窗口大小
}

type timeWindow struct {
	startTime int64
	endTime   int64
	offset    int // 数据块中的偏移
}

// NewTimeRangeIndex 创建新的时间范围索引
func NewTimeRangeIndex(blockSize int64) *TimeRangeIndex {
	return &TimeRangeIndex{
		windows:   make([]timeWindow, 0),
		blockSize: blockSize,
	}
}

// AddWindow 添加时间窗口
func (idx *TimeRangeIndex) AddWindow(start, end int64, offset int) {
	idx.windows = append(idx.windows, timeWindow{
		startTime: start,
		endTime:   end,
		offset:    offset,
	})
}

// FindBlocks 查找指定时间范围内的数据块
func (idx *TimeRangeIndex) FindBlocks(start, end int64) []int {
	var result []int
	for i, window := range idx.windows {
		if window.endTime >= start && window.startTime <= end {
			result = append(result, i)
		}
	}
	return result
}

// CircularBuffer 实现固定大小的环形缓冲区
type CircularBuffer struct {
	values   []DataPoint
	head     int
	size     int
	capacity int
	mu       sync.RWMutex
}

// NewCircularBuffer 创建新的环形缓冲区
func NewCircularBuffer(capacity int) *CircularBuffer {
	return &CircularBuffer{
		values:   make([]DataPoint, capacity),
		capacity: capacity,
	}
}

// Add 添加数据点到环形缓冲区
func (cb *CircularBuffer) Add(point DataPoint) {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	cb.values[cb.head] = point
	cb.head = (cb.head + 1) % cb.capacity
	if cb.size < cb.capacity {
		cb.size++
	}
}

// GetRecent 获取最近的n个数据点
func (cb *CircularBuffer) GetRecent(n int) []DataPoint {
	cb.mu.RLock()
	defer cb.mu.RUnlock()

	if n > cb.size {
		n = cb.size
	}

	result := make([]DataPoint, n)
	for i := 0; i < n; i++ {
		idx := (cb.head - i - 1 + cb.capacity) % cb.capacity
		if idx < 0 {
			idx += cb.capacity
		}
		result[i] = cb.values[idx]
	}

	return result
}

// Size 返回当前缓冲区中的数据点数量
func (cb *CircularBuffer) Size() int {
	cb.mu.RLock()
	defer cb.mu.RUnlock()
	return cb.size
}

// Capacity 返回缓冲区容量
func (cb *CircularBuffer) Capacity() int {
	return cb.capacity
}