从零实现一个TSDB（二）

二、实现启动函数和压缩算法和内存有序list

数据压缩算法，目前是两个算法，ZSTD和Snappy两种，也是直接调用的库函数，实现比较简单

package tsdb

import (
 "github.com/golang/snappy"
 "github.com/klauspost/compress/zstd"
)

type BytesCompressorType int8

const (
 // NoopBytesCompressor 默认不压缩
 NoopBytesCompressor BytesCompressorType = iota

 // ZSTDBytesCompressor 使用ZSTD压缩算法
 ZSTDBytesCompressor

 // SnappyBytesCompressor 使用snappy压缩算法
 SnappyBytesCompressor
)

// noopBytesCompressor 默认不压缩
type noopBytesCompressor struct{}
type zstdBytesCompressor struct{}
type snappyBytesCompressor struct{}

// BytesCompressor 数据压缩接口
type BytesCompressor interface {
 Compress(data []byte) []byte
 Decompress(data []byte) ([]byte, error)
}

//  默认压缩算法

func newNoopBytesCompressor() BytesCompressor {
 return &noopBytesCompressor{}
}

func (n *noopBytesCompressor) Compress(data []byte) []byte {
 return data
}

func (n *noopBytesCompressor) Decompress(data []byte) ([]byte, error) {
 return data, nil
}

// ZSTD 压缩算法

func newZSTDBytesCompressor() BytesCompressor {
 return &zstdBytesCompressor{}
}

func (n *zstdBytesCompressor) Compress(data []byte) []byte {
 encoder, _ := zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.SpeedFastest))
 return encoder.EncodeAll(data, make([]byte, 0, len(data)))
}

func (n *zstdBytesCompressor) Decompress(data []byte) ([]byte, error) {
 decoder, _ := zstd.NewReader(nil)
 return decoder.DecodeAll(data, nil)
}

// snappy 压缩算法

func newSnappyBytesCompressor() BytesCompressor {
 return &snappyBytesCompressor{}
}

func (n *snappyBytesCompressor) Compress(data []byte) []byte {
 return snappy.Encode(nil, data)
}

func (n *snappyBytesCompressor) Decompress(data []byte) ([]byte, error) {
 return snappy.Decode(nil, data)
}

然后进行打开创建一个db对象来进行相关操作

func OpenTSDB(opts ...Option) *TSDB {
 for _, opt := range opts {
  opt(defaultOpts)
 }

 db := &TSDB{
  segments: newSegmentList(),
  queue:    make(chan []*Row, defaultQueueSize),
 }

 // 加载文件
 db.loadFiles()

 worker := runtime.GOMAXPROCS(-1)
 db.ctx, db.cancel = context.WithCancel(context.Background())
 for i := 0; i < worker; i++ {
  // 刷盘
  go db.saveRows(db.ctx)
 }
 return db
}

可以看到 核心代码就是saveRows()和loadFiles()两个方法

先看loadFiles()

  func (db *TSDB) loadFiles() {
 mkdir(defaultOpts.dataPath)
 err := filepath.Walk(defaultOpts.dataPath, func(path string, info fs.FileInfo, err error) error {
  if err != nil {
   return fmt.Errorf("failed to traverse the dir: %s, err: %v", path, err)
  }
  // 文件后续都是默认以seg开头
  if !info.IsDir() || !strings.HasPrefix(info.Name(), "seg-") {
   return nil
  }

  files, err := ioutil.ReadDir(filepath.Join(defaultOpts.dataPath, info.Name()))
  if err != nil {
   return fmt.Errorf("failed to load the data storage, err: %v", err)
  }

  // 从磁盘加载出最近的segment数据进入内存
  nowDiskSegment := &diskSegment{}
  for _, file := range files {
   filename := filepath.Join(defaultOpts.dataPath, info.Name(), file.Name())
   if strings.EqualFold(file.Name(), "data") {
    mmapFile, err := OpenMMapFile(filename)
    if err != nil {
     return fmt.Errorf("failed to open mmap file %s, err: %v", filename, err)
    }
    nowDiskSegment.dataFd = mmapFile
    nowDiskSegment.dataFilename = filename
    nowDiskSegment.labelVs = newLabelValueList()
   }

   if strings.EqualFold(file.Name(), "meta") {
    data, err := ioutil.ReadFile(filename)
    if err != nil {
     return fmt.Errorf("failed to read file: %s, err: %v", filename, err)
    }
    // 构造meta文件数据格式
    desc := Desc{}
    if err = json.Unmarshal(data, &desc); err != nil {
     return fmt.Errorf("failed to json unmarshal meta file: %v", err)
    }
    nowDiskSegment.minTimestamp = desc.MinTimestamp
    nowDiskSegment.maxTimestamp = desc.MaxTimestamp
   }
  }
  db.segments.Add(nowDiskSegment)
  return nil
 })

 if err != nil {
  logrus.Error(err)
 }
}

loadFiles()中有这么几个函数

• OpenMMapFile
• segment.Add()

大致流程这样：

1. 先判断当前dir是否可以打开
2. 然后将其中的data数据文件通过mmap的方式加载进入内存segment
3. 在meta文件中获取当data数据文件的一些元数据信息
4. 将其添加进入内存的有序链表

实现MMAP方式读取文件

mmap.go

package tsdb

import (
 "errors"
 "fmt"
 "golang.org/x/sys/unix"
 "os"
)

// 实现通过mmap的方式读取文件

type MMapFile struct {
 file *os.File
 data []byte
}

func OpenMMapFile(path string) (mmapFile *MMapFile, err error) {
 file, err := os.Open(path)
 if err != nil {
  return nil, fmt.Errorf("open file error, path: %s", path)
 }
 defer func() {
  if err != nil {
   file.Close()
  }
 }()

 var size int
 info, err := file.Stat()
 if err != nil {
  return nil, fmt.Errorf("stat error, path: %s", path)
 }
 size = int(info.Size())
 data, err := syscallMmap(file, size)
 if err != nil {
  return nil, errors.New("mmap error")
 }
 return &MMapFile{
  file: file,
  data: data,
 }, nil
}

// syscallMmap unix创建内存映射
func syscallMmap(file *os.File, length int) ([]byte, error) {
 return unix.Mmap(int(file.Fd()), 0, length, unix.PROT_READ, unix.MAP_SHARED)
}

// syscallMunmap unix取消内存映射
func syscallMunmap(data []byte) (err error) {
 return unix.Munmap(data)
}

将内存segment添加进入有序列表List

segment.go

func (s *segmentList) Add(segment Segment) {
 s.mutex.Lock()
 defer s.mutex.Unlock()
 s.list.Add(segment.MinTs(), segment)
}

list.go

// List 排序链表结构
type List interface {
 Add(key int64, data interface{})
}

type node struct {
 height int
 key    int64
 value  interface{}
 left   *node
 right  *node
}

type avlTree struct {
 tree *node
}

func newTree() List {
 return &avlTree{
  &node{
   height: -2,
  },
 }
}

func (tree *avlTree) Add(key int64, value interface{}) {
 // todo 写avltree
}

当然 别忘了一步，那就是如何把文件的数据异构成一个segment，所以我们需要实现一个disksegment

disk.go

package tsdb

import (
 "github.com/sirupsen/logrus"
 "sync"
)

type diskSegment struct {
 dataFd       *MMapFile
 dataFilename string
 dir          string
 load         bool

 wait         sync.WaitGroup
 labelVs      *labelValueList
 indexMap     *diskIndexMap
 series       []metaSeries
 minTimestamp int64
 maxTimestamp int64

 seriesCount     int64
 dataPointsCount int64
}

func (ds *diskSegment) MinTs() int64 {
 return ds.minTimestamp
}

func (ds *diskSegment) InsertRows(_ []*Row) {
 logrus.Error("disk segments are not mutable")
}

内部有序数据结构：

• avlTree
• skiplist
• 红黑树
• art tree
• ...

目前系统为了简单，所以直接采用的avltree实现有序，当然这个肯定不可能在生产环境这么玩，后续会考虑其他数据结构的实现

list.go增删改查四个方法

• Add(key int64, data interface{})
• Remove(key int64) bool
• Range(start, end int64) Iter
• All() Iter

方法定义

type List interface {
 Add(key int64, data interface{})
 Remove(key int64) bool
 Range(start, end int64) Iter
 All() Iter
}

我们的key是通过timestamp来进行排序的，同时在调用range(start, end int64)或者 All()方法时，返回的是一个迭代器

所以我们定义一个迭代器对象

list.go

type Iter interface {
 Next() bool
 Value() interface{}
}

type iter struct {
 cursor int
 data   []interface{}
}

func (tree *avlTree) All() Iter {
 return tree.tree.values(0, math.MaxInt64)
}

func (it *iter) Next() bool {
 it.cursor++
 if len(it.data) > it.cursor {
  return true
 }
 return false
}

从上面可以看到，其实就是一个数组和一个游标cursor，临时保存当前遍历的数据segment而已

接下来就是实现方法

list.go的Add(key int64, value interface{})方法

func (tree *avlTree) Add(key int64, value interface{}) {
 // 插入insert avlTree
 tree.tree = insert(key, value, tree.tree)
}

其核心就是插入过程的具体实现和保持平衡的过程

list.go insert采用递归的方式去插入，而每次插入都需要判断一下当前tree是否balance

func insert(key int64, value interface{}, avlNode *node) *node {
 if avlNode == nil {
  return &node{
   key:   key,
   value: value,
  }
 }

 if avlNode.height == -2 {
  avlNode.key = key
  avlNode.value = value
  avlNode.height = 0
  return avlNode
 }

 diff := key - avlNode.key
 if diff > 0 {
  // 插入右子树
  avlNode.right = insert(key, value, avlNode.right)
 } else if diff < 0 {
  // 插入左子树
  avlNode.left = insert(key, value, avlNode.left)
 } else {
  avlNode.value = value
 }

 avlNode.keepBalance(key)
 avlNode.height = maxHeight(avlNode.left.height, avlNode.right.height) + 1
 return avlNode
}

list.go keepBalance方法

func (avlNode *node) keepBalance(key int64) *node {
 if avlNode.left.height-avlNode.right.height == 2 {
  if key-avlNode.left.key < 0 {
   avlNode = avlNode.rr()
  } else {
   avlNode = avlNode.lr()
  }
 } else if avlNode.right.height-avlNode.left.height == 2 {
  if avlNode.right.right.height > avlNode.left.height {
   avlNode = avlNode.rr()
  } else {
   avlNode = avlNode.rl()
  }
 }

 avlNode.height = maxHeight(avlNode.left.height, avlNode.right.height) + 1
 return avlNode
}

可以看到 每次平衡都得涉及到大量的左旋右旋操作。

// rr 插入节点在失衡节点的左子树的左子树中，需要右旋
func (avlNode *node) rr() *node {
 next := avlNode.left
 avlNode.left = next.right
 next.right = avlNode

 next.height = maxHeight(next.left.height, next.right.height) + 1
 avlNode.height = maxHeight(avlNode.left.height, avlNode.right.height) + 1
 return next
}

// lr 插入节点在失衡节点的左子树的右子树中，先左旋后右旋
func (avlNode *node) lr() *node {
 avlNode.left = avlNode.left.ll()
 return avlNode.rr()
}

// ll 插入的节点在失衡节点的右子树的右子树中，左旋
func (avlNode *node) ll() *node {
 next := avlNode.right
 avlNode.right = next.left
 next.left = avlNode

 next.height = maxHeight(next.left.height, next.right.height) + 1
 avlNode.height = maxHeight(avlNode.left.height, avlNode.right.height)
 return next
}

// rl 插入的节点在失衡节点的右子树的左子树中，先右旋后左旋
func (avlNode *node) rl() *node {
 avlNode.right = avlNode.right.rr()
 return avlNode.ll()
}

以上就是list.go的Add方法

接下来看看Remove方法

remove分成三步

1. 查找该key是否在该二叉树中
2. 将该key删除
3. 删除完如果树不是null，则需要进行平衡操作

find()函数查找该key是否在树中

func (avlNode *node) find(key int64) bool {
 if avlNode == nil {
  return false
 }
 diff := key - avlNode.key
 if diff > 0 {
  return avlNode.right.find(key)
 } else if diff < 0 {
  return avlNode.left.find(key)
 } else {
  return true
 }
}

delete()方法将该key从树中删除并进行平衡操作

func (avlNode *node) delete(key int64) *node {
 if avlNode == nil {
  return avlNode
 }

 diff := key - avlNode.key
 if diff > 0 {
  avlNode.right = avlNode.right.delete(key)
 } else if diff < 0 {
  avlNode.left = avlNode.left.delete(key)
 } else {
  if avlNode.left != nil && avlNode.right != nil {
   rightNode := avlNode.right.minNode()
   avlNode.key = rightNode.key
   avlNode.value = rightNode.value
   avlNode.right = avlNode.right.delete(avlNode.key)
  } else if avlNode.left != nil {
   avlNode = avlNode.left
  } else {
   avlNode = avlNode.right
  }
 }

 if avlNode != nil {
  avlNode.height = maxHeight(avlNode.left.height, avlNode.right.height) + 1
  avlNode = avlNode.keepBalance(key)
 }
 return avlNode
}

然后接下来就是All()和Range()方法的实现，其实这俩比较简单，就是遍历一遍二叉树将结果保存在一个数组中返回而已

func (tree *avlTree) Range(start, end int64) Iter {
 return tree.tree.values(start, end)
}

func (tree *avlTree) All() Iter {
 return tree.tree.values(0, math.MaxInt64)
}

核心就是values()方法，获取所有满足条件的值

func (avlNode *node) values(start, end int64) Iter {
 item := &iter{
  data: []interface{}{
   nil,
  },
 }
 item.data = list(item.data, start, end, avlNode)
 return item
}

func list(values []interface{}, start, end int64, avlNode *node) []interface{} {
 if avlNode != nil {
  values = list(values, start, end, avlNode.left)
  if avlNode.key >= start && avlNode.key <= end {
   values = append(values, avlNode.value)
  }
  values = list(values, start, end, avlNode.right)
 }
 return values
}

源码同步更新

github：https://github.com/azhsmesos/tsdb

本文由 mdnice 多平台发布