<|start_filename|>arenaskl/arena.go<|end_filename|>
/*
 * Copyright 2017 <NAME>, Inc. and Contributors
 * Modifications copyright (C) 2017 <NAME> and Contributors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package arenaskl

import (
	"math"
	"sync/atomic"
	"unsafe"
)

// Align requested alignment.
type Align uint8

const (
	// Align1 align 0.
	Align1 = 0
	// Align8 align 7
	Align8 = 7
)

// Arena should be lock-free.
type Arena struct {
	n   uint64
	buf []byte
}

// NewArena allocates a new arena of the specified size and returns it.
func NewArena(size uint32) *Arena {
	// Don't store data at position 0 in order to reserve offset=0 as a kind
	// of nil pointer.
	out := &Arena{
		n:   1,
		buf: make([]byte, size),
	}
	return out
}

// Size size have been allocated.
func (a *Arena) Size() uint32 {
	s := atomic.LoadUint64(&a.n)
	if s > math.MaxUint32 {
		// Saturate at MaxUint32.
		return math.MaxUint32
	}
	return uint32(s)
}

// Cap capacity of arena buffer.
func (a *Arena) Cap() uint32 {
	return uint32(len(a.buf))
}

// Reset reset size of arena.
func (a *Arena) Reset() {
	atomic.StoreUint64(&a.n, 1)
}

// Alloc allocate memory buffer by the given size.
func (a *Arena) Alloc(size, overflow uint32, align Align) (uint32, error) {
	// Verify that the arena isn't already full.
	origSize := atomic.LoadUint64(&a.n)
	if int(origSize) > len(a.buf) {
		a.growBufSize(origSize - uint64(len(a.buf)))
	}

	// Pad the allocation with enough bytes to ensure the requested alignment.
	padded := size + uint32(align)

	// Use 64-bit arithmetic to protect against overflow.
	newSize := atomic.AddUint64(&a.n, uint64(padded))
	if int(newSize)+int(overflow) > len(a.buf) {
		a.growBufSize(uint64(padded + overflow))
	}

	// Return the aligned offset.
	offset := (uint32(newSize) - padded + uint32(align)) & ^uint32(align)
	return offset, nil
}

// GetBytes returns byte slice at offset. The given size should be just the value
// size and should NOT include the meta bytes.
func (a *Arena) GetBytes(offset uint32, size uint32) []byte {
	if offset == 0 {
		return nil
	}

	if offset >= a.Cap() || offset+size >= a.Cap() {
		return nil
	}

	return a.buf[offset : offset+size]
}

// GetPointer get pointer at offset.
func (a *Arena) GetPointer(offset uint32) unsafe.Pointer {
	if offset == 0 {
		return nil
	}
	return unsafe.Pointer(&a.buf[offset])
}

// GetPointerOffset get the pointer offset in buffer.
func (a *Arena) GetPointerOffset(ptr unsafe.Pointer) uint32 {
	if ptr == nil {
		return 0
	}
	return uint32(uintptr(ptr) - uintptr(unsafe.Pointer(&a.buf[0])))
}

func (a *Arena) growBufSize(growBy uint64) {
	newBuf := make([]byte, uint64(len(a.buf))+growBy)
	copy(newBuf, a.buf)
	a.buf = newBuf
}


<|start_filename|>index/bptree.go<|end_filename|>
package index

import (
	"github.com/flower-corp/lotusdb/logger"
	"strings"
	"time"

	"go.etcd.io/bbolt"
)

const (
	defaultBatchSize = 100000
	fillPercent      = 0.9
)

// BPTreeOptions options for creating a new bptree.
type BPTreeOptions struct {
	// DirPath path to store index data.
	DirPath string
	// IndexType bptree(bolt).
	IndexType IndexerType
	// ColumnFamilyName db column family name, must be unique.
	ColumnFamilyName string
	// BucketName usually the the same as column family name, must be unique.
	BucketName []byte
	// BatchSize flush batch size.
	BatchSize int
	// DiscardChn if values in indexer are changed, the older value will be send to the DiscardChn.
	// Values will be handled in discard.go/listenUpdates().
	DiscardChn chan [][]byte
}

// BPTree is a standard b+tree used to store index data.
type BPTree struct {
	opts       BPTreeOptions
	db         *bbolt.DB
	discardChn chan [][]byte
}

// SetType self-explanatory.
func (bo *BPTreeOptions) SetType(typ IndexerType) {
	bo.IndexType = typ
}

// SetColumnFamilyName self-explanatory.
func (bo *BPTreeOptions) SetColumnFamilyName(cfName string) {
	bo.ColumnFamilyName = cfName
}

// SetDirPath self-explanatory.
func (bo *BPTreeOptions) SetDirPath(dirPath string) {
	bo.DirPath = dirPath
}

// GetType self-explanatory.
func (bo *BPTreeOptions) GetType() IndexerType {
	return bo.IndexType
}

// GetColumnFamilyName self-explanatory.
func (bo *BPTreeOptions) GetColumnFamilyName() string {
	return bo.ColumnFamilyName
}

// GetDirPath self-explanatory.
func (bo *BPTreeOptions) GetDirPath() string {
	return bo.DirPath
}

// NewBPTree create a boltdb instance.
// A file can only be opened once. if not, file lock competition will occur.
func NewBPTree(opt BPTreeOptions) (*BPTree, error) {
	if err := checkBPTreeOptions(opt); err != nil {
		return nil, err
	}

	// open metadatadb and db
	path := opt.DirPath + separator + strings.ToUpper(opt.GetColumnFamilyName())
	db, err := bbolt.Open(path+indexFileSuffixName, 0600, &bbolt.Options{
		Timeout:         1 * time.Second,
		NoSync:          true,
		InitialMmapSize: 1024,
		FreelistType:    bbolt.FreelistMapType,
	})
	if err != nil {
		return nil, err
	}

	tx, err := db.Begin(true)
	if err != nil {
		return nil, err
	}

	// cas create bucket
	if _, err := tx.CreateBucketIfNotExists(opt.BucketName); err != nil {
		return nil, err
	}
	// commit operation
	if err := tx.Commit(); err != nil {
		return nil, err
	}

	b := &BPTree{db: db, opts: opt, discardChn: opt.DiscardChn}
	return b, nil
}

// Put method starts a transaction.
// This method writes kv according to the bucket, and creates it if the bucket name does not exist.
func (b *BPTree) Put(key, value []byte) (err error) {
	var tx *bbolt.Tx
	if tx, err = b.db.Begin(true); err != nil {
		return
	}
	bucket := tx.Bucket(b.opts.BucketName)
	if _, err = bucket.Put(key, value); err != nil {
		_ = tx.Rollback()
		return
	}
	return tx.Commit()
}

// PutBatch is used for batch writing scenarios.
// The offset marks the transaction write position of the current batch.
// If this function fails during execution, we can write again from the offset position.
// If offset == len(kv) - 1 , all writes are successful.
func (b *BPTree) PutBatch(nodes []*IndexerNode, opts WriteOptions) (offset int, err error) {
	batchLoopNum := len(nodes) / b.opts.BatchSize
	if len(nodes)%b.opts.BatchSize > 0 {
		batchLoopNum++
	}

	batchlimit := b.opts.BatchSize
	for batchIdx := 0; batchIdx < batchLoopNum; batchIdx++ {
		offset = batchIdx * batchlimit
		tx, err := b.db.Begin(true)
		if err != nil {
			return offset, err
		}

		bucket := tx.Bucket(b.opts.BucketName)
		bucket.FillPercent = fillPercent
		var oldValues [][]byte
		for itemIdx := offset; itemIdx < offset+b.opts.BatchSize; itemIdx++ {
			if itemIdx >= len(nodes) {
				break
			}
			meta := EncodeMeta(nodes[itemIdx].Meta)
			if oldVal, err := bucket.Put(nodes[itemIdx].Key, meta); err != nil {
				_ = tx.Rollback()
				return offset, err
			} else if len(oldVal) > 0 && opts.SendDiscard {
				oldValues = append(oldValues, oldVal)
			}
		}
		if err := tx.Commit(); err != nil {
			return offset, err
		}
		if opts.SendDiscard {
			b.sendDiscard(oldValues)
		}
	}
	return len(nodes) - 1, nil
}

// DeleteBatch delete data in batch.
func (b *BPTree) DeleteBatch(keys [][]byte, opts WriteOptions) error {
	batchLoopNum := len(keys) / b.opts.BatchSize
	if len(keys)%b.opts.BatchSize > 0 {
		batchLoopNum++
	}

	batchlimit := b.opts.BatchSize
	for batchIdx := 0; batchIdx < batchLoopNum; batchIdx++ {
		offset := batchIdx * batchlimit
		tx, err := b.db.Begin(true)
		if err != nil {
			return err
		}
		bucket := tx.Bucket(b.opts.BucketName)
		bucket.FillPercent = fillPercent
		var oldValues [][]byte
		for itemIdx := offset; itemIdx < offset+b.opts.BatchSize; itemIdx++ {
			if itemIdx >= len(keys) {
				break
			}
			if oldVal, err := bucket.Delete(keys[itemIdx]); err != nil {
				_ = tx.Rollback()
				return err
			} else if len(oldVal) > 0 && opts.SendDiscard {
				oldValues = append(oldValues, oldVal)
			}
		}
		if err := tx.Commit(); err != nil {
			return err
		}
		if opts.SendDiscard {
			b.sendDiscard(oldValues)
		}
	}
	return nil
}

// Delete a specified key from indexer.
func (b *BPTree) Delete(key []byte) error {
	return b.db.Update(func(tx *bbolt.Tx) error {
		_, err := tx.Bucket(b.opts.BucketName).Delete(key)
		return err
	})
}

// Get reads the value from the bucket with key.
func (b *BPTree) Get(key []byte) (*IndexerMeta, error) {
	tx, err := b.db.Begin(false)
	if err != nil {
		return nil, err
	}
	defer func() {
		_ = tx.Rollback()
	}()

	buf := tx.Bucket(b.opts.BucketName).Get(key)
	if buf == nil {
		return nil, nil
	}
	return DecodeMeta(buf), nil
}

// Sync executes fdatasync() against the database file handle.
func (b *BPTree) Sync() error {
	return b.db.Sync()
}

// Close close bolt db.
func (b *BPTree) Close() error {
	return b.db.Close()
}

func (b *BPTree) sendDiscard(values [][]byte) {
	if len(values) > 0 {
		select {
		case b.discardChn <- values:
		default:
			logger.Warn("send to discard chan fail")
		}
	}
}

func checkBPTreeOptions(opt BPTreeOptions) error {
	if opt.ColumnFamilyName == "" {
		return ErrColumnFamilyNameNil
	}

	if opt.DirPath == "" {
		return ErrDirPathNil
	}

	if opt.BucketName == nil || len(opt.BucketName) == 0 {
		return ErrBucketNameNil
	}

	if opt.BatchSize < defaultBatchSize {
		opt.BatchSize = defaultBatchSize
	}
	return nil
}


<|start_filename|>vlog.go<|end_filename|>
package lotusdb

import (
	"errors"
	"fmt"
	"github.com/flower-corp/lotusdb/index"
	"github.com/flower-corp/lotusdb/logfile"
	"github.com/flower-corp/lotusdb/logger"
	"io"
	"io/ioutil"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
)

var (
	// ErrActiveLogFileNil active log file not exists.
	ErrActiveLogFileNil = errors.New("active log file not exists")

	// ErrLogFileNil log file not exists.
	ErrLogFileNil = errors.New("log file %d not exists")
)

const vlogDiscardName = "VLOG_DISCARD"

type (
	// valueLog value log is named after the concept in Wisckey paper(https://www.usenix.org/system/files/conference/fast16/fast16-papers-lu.pdf).
	// Values will be stored in value log if its size exceed ValueThreshold in options.
	valueLog struct {
		sync.RWMutex
		opt           vlogOptions
		activeLogFile *logfile.LogFile            // current active log file for writing.
		logFiles      map[uint32]*logfile.LogFile // all log files. Must hold the mutex before modify it.
		cf            *ColumnFamily
		discard       *discard
	}

	// valuePos value position.
	valuePos struct {
		Fid    uint32
		Offset int64
	}

	vlogOptions struct {
		path       string
		blockSize  int64
		ioType     logfile.IOType
		gcRatio    float64
		gcInterval time.Duration
	}
)

// openValueLog create a new value log file.
func openValueLog(opt vlogOptions) (*valueLog, error) {
	fileInfos, err := ioutil.ReadDir(opt.path)
	if err != nil {
		return nil, err
	}

	var fids []uint32
	for _, file := range fileInfos {
		if strings.HasSuffix(file.Name(), logfile.VLogSuffixName) {
			splitNames := strings.Split(file.Name(), ".")
			fid, err := strconv.Atoi(splitNames[0])
			if err != nil {
				return nil, err
			}
			fids = append(fids, uint32(fid))
		}
	}

	// load in order.
	sort.Slice(fids, func(i, j int) bool {
		return fids[i] < fids[j]
	})
	if len(fids) == 0 {
		fids = append(fids, logfile.InitialLogFileId)
	}

	// open discard file.
	discard, err := newDiscard(opt.path, vlogDiscardName)
	if err != nil {
		return nil, err
	}

	// open active log file only.
	logFile, err := logfile.OpenLogFile(opt.path, fids[len(fids)-1], opt.blockSize, logfile.ValueLog, opt.ioType)
	// set total size in discard file, skip it if exist.
	discard.setTotal(fids[len(fids)-1], uint32(opt.blockSize))
	if err != nil {
		return nil, err
	}

	vlog := &valueLog{
		opt:           opt,
		activeLogFile: logFile,
		logFiles:      make(map[uint32]*logfile.LogFile),
		discard:       discard,
	}

	// load other log files when reading from it.
	for i := 0; i < len(fids)-1; i++ {
		vlog.logFiles[fids[i]] = &logfile.LogFile{Fid: fids[i]}
	}

	if err := vlog.setLogFileState(); err != nil {
		return nil, err
	}
	go vlog.handleCompaction()
	return vlog, nil
}

// Read a VLogEntry from a specified vlog file at offset, returns an error, if any.
// If reading from a non-active log file, and the specified file is not open, then we will open it and set it into logFiles.
func (vlog *valueLog) Read(fid uint32, offset int64) (*logfile.LogEntry, error) {
	var logFile *logfile.LogFile
	if fid == vlog.activeLogFile.Fid {
		logFile = vlog.activeLogFile
	} else {
		vlog.RLock()
		logFile = vlog.logFiles[fid]
		if logFile != nil && logFile.IoSelector == nil {
			opt := vlog.opt
			lf, err := logfile.OpenLogFile(opt.path, fid, opt.blockSize, logfile.ValueLog, opt.ioType)
			if err != nil {
				vlog.RUnlock()
				return nil, err
			}
			vlog.logFiles[fid] = lf
			logFile = lf
		}
		vlog.RUnlock()
	}
	if logFile == nil {
		return nil, fmt.Errorf(ErrLogFileNil.Error(), fid)
	}

	entry, _, err := logFile.ReadLogEntry(offset)
	if err == logfile.ErrEndOfEntry {
		return &logfile.LogEntry{}, nil
	}
	return entry, err
}

// Write new VLogEntry to value log file.
// If the active log file is full, it will be closed and a new active file will be created to replace it.
func (vlog *valueLog) Write(ent *logfile.LogEntry) (*valuePos, int, error) {
	vlog.Lock()
	defer vlog.Unlock()
	buf, eSize := logfile.EncodeEntry(ent)
	// if active is reach to thereshold, close it and open a new one.
	if vlog.activeLogFile.WriteAt+int64(eSize) >= vlog.opt.blockSize {
		if err := vlog.Sync(); err != nil {
			return nil, 0, err
		}
		vlog.logFiles[vlog.activeLogFile.Fid] = vlog.activeLogFile

		logFile, err := vlog.createLogFile()
		if err != nil {
			return nil, 0, err
		}
		vlog.activeLogFile = logFile
	}

	err := vlog.activeLogFile.Write(buf)
	if err != nil {
		return nil, 0, err
	}

	writeAt := atomic.LoadInt64(&vlog.activeLogFile.WriteAt)
	return &valuePos{
		Fid:    vlog.activeLogFile.Fid,
		Offset: writeAt - int64(eSize),
	}, eSize, nil
}

// Sync only for the active log file.
func (vlog *valueLog) Sync() error {
	if vlog.activeLogFile == nil {
		return ErrActiveLogFileNil
	}

	vlog.activeLogFile.Lock()
	defer vlog.activeLogFile.Unlock()
	return vlog.activeLogFile.Sync()
}

// Close only for the active log file.
func (vlog *valueLog) Close() error {
	if vlog.activeLogFile == nil {
		return ErrActiveLogFileNil
	}

	vlog.activeLogFile.Lock()
	defer vlog.activeLogFile.Unlock()
	return vlog.activeLogFile.Close()
}

func (vlog *valueLog) createLogFile() (*logfile.LogFile, error) {
	opt := vlog.opt
	fid := vlog.activeLogFile.Fid
	logFile, err := logfile.OpenLogFile(opt.path, fid+1, opt.blockSize, logfile.ValueLog, opt.ioType)
	if err != nil {
		return nil, err
	}
	vlog.discard.setTotal(fid+1, uint32(opt.blockSize))
	return logFile, nil
}

func (vlog *valueLog) setLogFileState() error {
	if vlog.activeLogFile == nil {
		return ErrActiveLogFileNil
	}
	var offset int64 = 0
	for {
		if _, size, err := vlog.activeLogFile.ReadLogEntry(offset); err == nil {
			offset += size
			// No need to use atomic updates.
			// This function is only be executed in one goroutine at startup.
			vlog.activeLogFile.WriteAt += size
		} else {
			if err == io.EOF || err == logfile.ErrEndOfEntry {
				break
			}
			return err
		}
	}
	// if active file`s capacity is nearly close to block size, open a new active file.
	if vlog.activeLogFile.WriteAt+logfile.MaxHeaderSize >= vlog.opt.blockSize {
		vlog.logFiles[vlog.activeLogFile.Fid] = vlog.activeLogFile
		logFile, err := vlog.createLogFile()
		if err != nil {
			return err
		}
		vlog.activeLogFile = logFile
	}
	return nil
}

func (vlog *valueLog) getActiveFid() uint32 {
	var fid uint32
	vlog.Lock()
	if vlog.activeLogFile != nil {
		fid = vlog.activeLogFile.Fid
	}
	vlog.Unlock()
	return fid
}

func (vlog *valueLog) getLogFile(fid uint32) *logfile.LogFile {
	vlog.Lock()
	defer vlog.Unlock()
	return vlog.logFiles[fid]
}

func (vlog *valueLog) handleCompaction() {
	if vlog.opt.gcInterval <= 0 {
		return
	}

	quitSig := make(chan os.Signal, 1)
	signal.Notify(quitSig, os.Interrupt, os.Kill, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
	ticker := time.NewTicker(vlog.opt.gcInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			if err := vlog.compact(); err != nil {
				logger.Errorf("value log compaction err: %+v", err)
			}
		case <-quitSig:
		case <-vlog.cf.closedC:
			return
		}
	}
}

func (vlog *valueLog) compact() error {
	rewrite := func(file *logfile.LogFile) error {
		vlog.cf.flushLock.Lock()
		defer vlog.cf.flushLock.Unlock()

		var offset int64
		var validEntries []*logfile.LogEntry
		ts := time.Now().Unix()
		for {
			entry, sz, err := file.ReadLogEntry(offset)
			if err != nil {
				if err == io.EOF || err == logfile.ErrEndOfEntry {
					break
				}
				return err
			}
			var eoff = offset
			offset += sz
			indexMeta, err := vlog.cf.indexer.Get(entry.Key)
			if err != nil {
				return err
			}
			if indexMeta == nil {
				continue
			}
			// if value is already stored in indexer, the value in vlog must be old.
			if len(indexMeta.Value) != 0 {
				continue
			}
			if indexMeta.Fid == file.Fid && indexMeta.Offset == eoff {
				validEntries = append(validEntries, entry)
			}
		}

		var nodes []*index.IndexerNode
		var deletedKeys [][]byte
		// rewrite valid log entries.
		for _, e := range validEntries {
			if e.ExpiredAt != 0 && e.ExpiredAt <= ts {
				deletedKeys = append(deletedKeys, e.Key)
				continue
			}

			valuePos, esize, err := vlog.Write(e)
			if err != nil {
				return err
			}
			nodes = append(nodes, &index.IndexerNode{
				Key: e.Key,
				Meta: &index.IndexerMeta{
					Fid:       valuePos.Fid,
					Offset:    valuePos.Offset,
					EntrySize: esize,
				},
			})
		}
		writeOpts := index.WriteOptions{SendDiscard: false}
		if _, err := vlog.cf.indexer.PutBatch(nodes, writeOpts); err != nil {
			return err
		}
		if len(deletedKeys) > 0 {
			if err := vlog.cf.indexer.DeleteBatch(deletedKeys, writeOpts); err != nil {
				return err
			}
		}
		if err := vlog.cf.indexer.Sync(); err != nil {
			return err
		}
		return nil
	}

	opt := vlog.opt
	activeFid := vlog.getActiveFid()
	ccl, err := vlog.discard.getCCL(activeFid, opt.gcRatio)
	if err != nil {
		return err
	}

	for _, fid := range ccl {
		lf := vlog.getLogFile(fid)
		if lf == nil || lf.IoSelector == nil {
			file, err := logfile.OpenLogFile(opt.path, fid, opt.blockSize, logfile.ValueLog, opt.ioType)
			if err != nil {
				return err
			}
			lf = file
		}

		if err = rewrite(lf); err != nil {
			logger.Warnf("compact rewrite err: %+v", err)
			return err
		}
		// clear discard state.
		vlog.discard.clear(fid)

		// delete older vlog file.
		vlog.Lock()
		if _, ok := vlog.logFiles[fid]; ok {
			delete(vlog.logFiles, fid)
		}
		vlog.Unlock()
		if err = lf.Delete(); err != nil {
			return err
		}
	}
	return nil
}
