memtable.go

package lotusdb

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"math"
	"os"
	"sort"
	"sync"

	"github.com/bwmarrin/snowflake"

	arenaskl "github.com/dgraph-io/badger/v4/skl"
	"github.com/dgraph-io/badger/v4/y"
	"github.com/rosedblabs/wal"
)

const (
	// the wal file name format is .SEG.%d
	// %d is the unique id of the memtable, used to generate wal file name
	// for example, the wal file name of memtable with id 1 is .SEG.1.
	walFileExt     = ".SEG.%d"
	initialTableID = 1
)

type (
	// memtable is an in-memory data structure holding data before they are flushed into index and value log.
	// Currently, the only supported data structure is skip list, see github.com/dgraph-io/badger/v4/skl.
	//
	// New writes always insert data to memtable, and reads has query from memtable
	// before reading from index and value log, because memtable`s data is newer.
	//
	// Once a memtable is full(memtable has its threshold, see MemtableSize in options),
	// it becomes immutable and replaced by a new memtable.
	//
	// A background goroutine will flush the content of memtable into index and vlog,
	// after that the memtable can be deleted.
	memtable struct {
		mu      sync.RWMutex
		wal     *wal.WAL           // write ahead log for the memtable
		skl     *arenaskl.Skiplist // in-memory skip list
		options memtableOptions
	}

	// memtableOptions represents the configuration options for a memtable.
	memtableOptions struct {
		dirPath         string // where write ahead log wal file is stored
		tableID         uint32 // unique id of the memtable, used to generate wal file name
		memSize         uint32 // max size of the memtable
		walBytesPerSync uint32 // flush wal file to disk throughput BytesPerSync parameter
		walSync         bool   // WAL flush immediately after each writing
	}
)

// find the wal file of the memtable with the specified id
// a wal is associated with a memtable, so the wal file name is generated by the memtable id
// for example, the wal file name of memtable with id 1 is .SEG.1.
func openAllMemtables(options Options) ([]*memtable, error) {
	entries, err := os.ReadDir(options.DirPath)
	if err != nil {
		return nil, err
	}

	// get all memtable ids
	var tableIDs []int
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		var id int
		var prefix int
		_, err = fmt.Sscanf(entry.Name(), "%d"+walFileExt, &prefix, &id)
		if err != nil {
			continue
		}
		tableIDs = append(tableIDs, id)
	}

	if len(tableIDs) == 0 {
		tableIDs = append(tableIDs, initialTableID)
	}
	sort.Ints(tableIDs)
	tables := make([]*memtable, len(tableIDs))
	for i, table := range tableIDs {
		table, errOpenMemtable := openMemtable(memtableOptions{
			dirPath:         options.DirPath,
			tableID:         uint32(table),
			memSize:         options.MemtableSize,
			walSync:         options.Sync,
			walBytesPerSync: options.BytesPerSync,
		})
		if errOpenMemtable != nil {
			return nil, errOpenMemtable
		}
		tables[i] = table
	}

	return tables, nil
}

// memtable holds a wal(write ahead log), so when opening a memtable,
// actually it open the corresponding wal file.
// and load all entries from wal to rebuild the content of the skip list.
func openMemtable(options memtableOptions) (*memtable, error) {
	// init skip list
	//nolint:gomnd // default size
	skl := arenaskl.NewSkiplist(int64(float64(options.memSize) * 1.5))
	table := &memtable{options: options, skl: skl}

	// open the Write Ahead Log file
	walFile, err := wal.Open(wal.Options{
		DirPath:        options.dirPath,
		SegmentSize:    math.MaxInt, // no limit, guarantee that a wal file only contains one segment file
		SegmentFileExt: fmt.Sprintf(walFileExt, options.tableID),
		Sync:           options.walSync,
		BytesPerSync:   options.walBytesPerSync,
	})
	if err != nil {
		return nil, err
	}
	table.wal = walFile

	indexRecords := make(map[uint64][]*LogRecord)
	// now we get the opened wal file, we need to load all entries
	// from wal to rebuild the content of the skip list
	reader := table.wal.NewReader()
	for {
		chunk, _, errNext := reader.Next()
		if errNext != nil {
			if errors.Is(errNext, io.EOF) {
				break
			}
			return nil, errNext
		}
		record := decodeLogRecord(chunk)
		if record.Type == LogRecordBatchFinished {
			batchID, errParseBytes := snowflake.ParseBytes(record.Key)
			if errParseBytes != nil {
				return nil, errParseBytes
			}
			for _, idxRecord := range indexRecords[uint64(batchID)] {
				table.skl.Put(y.KeyWithTs(idxRecord.Key, 0),
					y.ValueStruct{Value: idxRecord.Value, Meta: idxRecord.Type})
			}
			delete(indexRecords, uint64(batchID))
		} else {
			indexRecords[record.BatchID] = append(indexRecords[record.BatchID], record)
		}
	}

	// open and read wal file successfully, return the memtable
	return table, nil
}

// putBatch writes a batch of entries to memtable.
func (mt *memtable) putBatch(pendingWrites map[string]*LogRecord,
	batchID snowflake.ID, options WriteOptions) error {
	// if wal is not disabled, write to wal first to ensure durability and atomicity
	if !options.DisableWal {
		// add record to wal.pendingWrites
		for _, record := range pendingWrites {
			record.BatchID = uint64(batchID)
			encRecord := encodeLogRecord(record)
			mt.wal.PendingWrites(encRecord)
		}

		// add a record to indicate the end of the batch
		endRecord := encodeLogRecord(&LogRecord{
			Key:  batchID.Bytes(),
			Type: LogRecordBatchFinished,
		})
		mt.wal.PendingWrites(endRecord)

		// write wal.pendingWrites
		if _, err := mt.wal.WriteAll(); err != nil {
			return err
		}
		// flush wal if necessary
		if options.Sync && !mt.options.walSync {
			if err := mt.wal.Sync(); err != nil {
				return err
			}
		}
	}

	mt.mu.Lock()
	// write to in-memory skip list
	for key, record := range pendingWrites {
		mt.skl.Put(y.KeyWithTs([]byte(key), 0), y.ValueStruct{Value: record.Value, Meta: record.Type})
	}
	mt.mu.Unlock()

	return nil
}

// get value from memtable
// if the specified key is marked as deleted, a true bool value is returned.
func (mt *memtable) get(key []byte) (bool, []byte) {
	mt.mu.RLock()
	defer mt.mu.RUnlock()

	valueStruct := mt.skl.Get(y.KeyWithTs(key, 0))
	deleted := valueStruct.Meta == LogRecordDeleted
	return deleted, valueStruct.Value
}

func (mt *memtable) isFull() bool {
	return mt.skl.MemSize() >= int64(mt.options.memSize)
}

func (mt *memtable) deleteWAl() error {
	if mt.wal != nil {
		return mt.wal.Delete()
	}
	return nil
}

func (mt *memtable) close() error {
	if mt.wal != nil {
		return mt.wal.Close()
	}
	return nil
}

func (mt *memtable) sync() error {
	if mt.wal != nil {
		return mt.wal.Sync()
	}
	return nil
}

// memtableIterator implement baseIterator.
type memtableIterator struct {
	options IteratorOptions
	iter    *arenaskl.UniIterator
}

func newMemtableIterator(options IteratorOptions, memtable *memtable) *memtableIterator {
	return &memtableIterator{
		options: options,
		iter:    memtable.skl.NewUniIterator(options.Reverse),
	}
}

// Rewind seek the first key in the iterator.
func (mi *memtableIterator) Rewind() {
	mi.iter.Rewind()
	if len(mi.options.Prefix) == 0 {
		return
	}
	// prefix scan
	for mi.iter.Valid() && !bytes.HasPrefix(mi.iter.Key(), mi.options.Prefix) {
		mi.iter.Next()
	}
}

// Seek move the iterator to the key which is
// greater(less when reverse is true) than or equal to the specified key.
func (mi *memtableIterator) Seek(key []byte) {
	mi.iter.Seek(y.KeyWithTs(key, 0))
	if len(mi.options.Prefix) == 0 {
		return
	}
	// prefix scan
	for mi.Valid() && !bytes.HasPrefix(mi.Key(), mi.options.Prefix) {
		mi.Next()
	}
}

// Next moves the iterator to the next key.
func (mi *memtableIterator) Next() {
	mi.iter.Next()
	if len(mi.options.Prefix) == 0 {
		return
	}
	// prefix scan
	for mi.iter.Valid() && !bytes.HasPrefix(mi.iter.Key(), mi.options.Prefix) {
		mi.iter.Next()
	}
}

// Key get the current key.
func (mi *memtableIterator) Key() []byte {
	return y.ParseKey(mi.iter.Key())
}

// Value get the current value.
func (mi *memtableIterator) Value() any {
	return mi.iter.Value()
}

// Valid returns whether the iterator is exhausted.
func (mi *memtableIterator) Valid() bool {
	return mi.iter.Valid()
}

// Close the iterator.
func (mi *memtableIterator) Close() error {
	return mi.iter.Close()
}