541 lines
15 KiB
Nim
541 lines
15 KiB
Nim
## LSM-Tree Storage Engine — core key-value store
|
|
import std/algorithm
|
|
import std/os
|
|
import std/hashes
|
|
import std/strutils
|
|
import std/tables
|
|
import std/monotimes
|
|
import std/streams
|
|
import std/locks
|
|
import bloom
|
|
import wal
|
|
import mmap
|
|
|
|
const
|
|
SSTableMagic* = 0x53535442'u32 # "SSTB"
|
|
SSTableVersion* = 1'u32
|
|
DefaultMemTableSize* = 4 * 1024 * 1024 # 4MB
|
|
DefaultBloomFpRate* = 0.01
|
|
|
|
type
|
|
Entry* = object
|
|
key*: string
|
|
value*: seq[byte]
|
|
timestamp*: uint64
|
|
deleted*: bool
|
|
|
|
MemTable* = object
|
|
entries: seq[Entry]
|
|
size: int
|
|
maxSize: int
|
|
|
|
SSTable* = object
|
|
id*: int
|
|
path*: string
|
|
index*: Table[string, int64]
|
|
bloom*: BloomFilter
|
|
level*: int
|
|
minKey*: string
|
|
maxKey*: string
|
|
entryCount*: int
|
|
mmapFile*: MmapFile
|
|
|
|
LSMTree* = ref object
|
|
dir*: string
|
|
memTable: MemTable
|
|
immutableMem: MemTable
|
|
sstables*: seq[SSTable]
|
|
wal: WriteAheadLog
|
|
memMaxSize: int
|
|
currentSeq: uint64
|
|
nextSSTableId: int
|
|
lock*: Lock
|
|
walLock*: Lock
|
|
|
|
proc newMemTable(maxSize: int = DefaultMemTableSize): MemTable =
|
|
MemTable(entries: @[], size: 0, maxSize: maxSize)
|
|
|
|
proc len*(mt: MemTable): int = mt.entries.len
|
|
|
|
proc put*(mt: var MemTable, key: string, value: seq[byte], timestamp: uint64, deleted: bool = false): bool =
|
|
let entrySize = key.len + value.len + 16
|
|
if mt.size + entrySize > mt.maxSize and mt.entries.len > 0:
|
|
return false
|
|
let entry = Entry(key: key, value: value, timestamp: timestamp, deleted: deleted)
|
|
let pos = mt.entries.lowerBound(entry, proc(a, b: Entry): int = cmp(a.key, b.key))
|
|
if pos < mt.entries.len and mt.entries[pos].key == key:
|
|
mt.entries[pos] = entry
|
|
else:
|
|
mt.entries.insert(entry, pos)
|
|
mt.size += entrySize
|
|
return true
|
|
|
|
proc get*(mt: MemTable, key: string): (bool, Entry) =
|
|
if mt.entries.len == 0:
|
|
return (false, Entry())
|
|
var lo = 0
|
|
var hi = mt.entries.len - 1
|
|
while lo <= hi:
|
|
let mid = (lo + hi) div 2
|
|
let c = cmp(mt.entries[mid].key, key)
|
|
if c == 0:
|
|
return (true, mt.entries[mid])
|
|
elif c < 0:
|
|
lo = mid + 1
|
|
else:
|
|
hi = mid - 1
|
|
return (false, Entry())
|
|
|
|
proc scan*(mt: MemTable, startKey, endKey: string): seq[Entry] =
|
|
result = @[]
|
|
for entry in mt.entries:
|
|
if entry.key >= startKey and entry.key <= endKey:
|
|
result.add(entry)
|
|
|
|
proc clear*(mt: var MemTable) =
|
|
mt.entries.setLen(0)
|
|
mt.size = 0
|
|
|
|
# ----------------------------------------------------------------------
|
|
# SSTable serialization format (native endianness):
|
|
# [Header] 28 bytes
|
|
# magic: uint32
|
|
# version: uint32
|
|
# entryCount: uint32
|
|
# indexOffset: uint64
|
|
# bloomOffset: uint64
|
|
#
|
|
# [Data Block]
|
|
# For each entry:
|
|
# keyLen: uint32
|
|
# key: bytes[keyLen]
|
|
# valueLen: uint32
|
|
# value: bytes[valueLen]
|
|
# timestamp: uint64
|
|
# deleted: uint8
|
|
#
|
|
# [Index Block]
|
|
# For each entry:
|
|
# keyLen: uint32
|
|
# key: bytes[keyLen]
|
|
# dataOffset: uint64 # offset of this entry in data block
|
|
#
|
|
# [Bloom Filter Block]
|
|
# bloomSize: uint32
|
|
# bloomData: bytes[bloomSize]
|
|
# ----------------------------------------------------------------------
|
|
|
|
proc writeSSTable*(entries: seq[Entry], path: string, level: int): SSTable =
|
|
let s = newFileStream(path, fmWrite)
|
|
if s.isNil:
|
|
raise newException(IOError, "Cannot create SSTable file: " & path)
|
|
|
|
# Write header
|
|
s.write(SSTableMagic)
|
|
s.write(SSTableVersion)
|
|
s.write(uint32(entries.len))
|
|
let indexOffsetPos = s.getPosition()
|
|
s.write(0'u64) # patched after data+bloom are written
|
|
let bloomOffsetPos = s.getPosition()
|
|
s.write(0'u64) # patched after data+bloom are written
|
|
|
|
# Write data block
|
|
var offsets = newSeq[(string, int64)](entries.len)
|
|
for i, entry in entries:
|
|
offsets[i] = (entry.key, int64(s.getPosition()))
|
|
s.write(uint32(entry.key.len))
|
|
s.write(entry.key)
|
|
s.write(uint32(entry.value.len))
|
|
if entry.value.len > 0:
|
|
s.writeData(addr entry.value[0], entry.value.len)
|
|
s.write(entry.timestamp)
|
|
s.write(if entry.deleted: 1'u8 else: 0'u8)
|
|
|
|
let indexOffset = uint64(s.getPosition())
|
|
|
|
# Write index block
|
|
for i, entry in entries:
|
|
s.write(uint32(entry.key.len))
|
|
s.write(entry.key)
|
|
s.write(uint64(offsets[i][1]))
|
|
|
|
let bloomOffset = uint64(s.getPosition())
|
|
|
|
# Write bloom filter
|
|
var bloom = newBloomFilter(max(entries.len * 10, 1000), DefaultBloomFpRate)
|
|
for entry in entries:
|
|
bloom.add(cast[seq[byte]](entry.key))
|
|
let bloomData = bloom.serialize()
|
|
s.write(uint32(bloomData.len))
|
|
if bloomData.len > 0:
|
|
s.writeData(addr bloomData[0], bloomData.len)
|
|
|
|
# Patch header with correct offsets before closing
|
|
s.setPosition(int(indexOffsetPos))
|
|
s.write(indexOffset)
|
|
s.setPosition(int(bloomOffsetPos))
|
|
s.write(bloomOffset)
|
|
s.close()
|
|
|
|
# Build in-memory index
|
|
var idxTable = initTable[string, int64]()
|
|
var minK = ""
|
|
var maxK = ""
|
|
for i, entry in entries:
|
|
idxTable[entry.key] = offsets[i][1]
|
|
if minK == "" or entry.key < minK: minK = entry.key
|
|
if maxK == "" or entry.key > maxK: maxK = entry.key
|
|
|
|
result = SSTable(
|
|
id: -1,
|
|
path: path,
|
|
index: idxTable,
|
|
bloom: bloom,
|
|
level: level,
|
|
minKey: minK,
|
|
maxKey: maxK,
|
|
entryCount: entries.len,
|
|
mmapFile: openMmap(path),
|
|
)
|
|
|
|
proc loadSSTable*(path: string): SSTable =
|
|
let mf = openMmap(path)
|
|
if mf.regions.len == 0:
|
|
raise newException(IOError, "Cannot mmap SSTable: " & path)
|
|
|
|
if mf.readUint32(0) != SSTableMagic:
|
|
raise newException(ValueError, "Invalid SSTable magic")
|
|
if mf.readUint32(4) != SSTableVersion:
|
|
raise newException(ValueError, "Unsupported SSTable version")
|
|
|
|
let entryCount = int(mf.readUint32(8))
|
|
let indexOffset = int(mf.readUint64(12))
|
|
let bloomOffset = int(mf.readUint64(20))
|
|
|
|
var idxTable = initTable[string, int64]()
|
|
var minK = ""
|
|
var maxK = ""
|
|
|
|
# Parse index block
|
|
var pos = indexOffset
|
|
for i in 0..<entryCount:
|
|
let keyLen = int(mf.readUint32(pos))
|
|
pos += 4
|
|
let key = mf.readString(pos, keyLen)
|
|
pos += keyLen
|
|
let dataOffset = int64(mf.readUint64(pos))
|
|
pos += 8
|
|
idxTable[key] = dataOffset
|
|
if minK == "" or key < minK: minK = key
|
|
if maxK == "" or key > maxK: maxK = key
|
|
|
|
# Parse bloom filter
|
|
var bloom = newBloomFilter(max(entryCount * 10, 1000), DefaultBloomFpRate)
|
|
pos = bloomOffset
|
|
let bloomSize = int(mf.readUint32(pos))
|
|
pos += 4
|
|
if bloomSize > 0 and pos + bloomSize <= mf.totalSize:
|
|
var bloomData = newSeq[byte](bloomSize)
|
|
copyMem(addr bloomData[0], unsafeAddr mf.regions[0].data[pos], bloomSize)
|
|
bloom.deserialize(bloomData)
|
|
|
|
result = SSTable(
|
|
id: -1,
|
|
path: path,
|
|
index: idxTable,
|
|
bloom: bloom,
|
|
level: 0,
|
|
minKey: minK,
|
|
maxKey: maxK,
|
|
entryCount: entryCount,
|
|
mmapFile: mf,
|
|
)
|
|
|
|
proc readSSTableEntry*(sst: SSTable, key: string): (bool, Entry) =
|
|
if key notin sst.index:
|
|
return (false, Entry())
|
|
|
|
let offset = int(sst.index[key])
|
|
let mf = sst.mmapFile
|
|
if mf.regions.len == 0:
|
|
return (false, Entry())
|
|
|
|
var pos = offset
|
|
if pos + 4 > mf.totalSize:
|
|
return (false, Entry())
|
|
let keyLen = int(mf.readUint32(pos))
|
|
pos += 4
|
|
if pos + keyLen > mf.totalSize:
|
|
return (false, Entry())
|
|
let readKey = mf.readString(pos, keyLen)
|
|
pos += keyLen
|
|
if readKey != key:
|
|
return (false, Entry())
|
|
|
|
if pos + 4 > mf.totalSize:
|
|
return (false, Entry())
|
|
let valueLen = int(mf.readUint32(pos))
|
|
pos += 4
|
|
if pos + valueLen + 8 + 1 > mf.totalSize:
|
|
return (false, Entry())
|
|
|
|
var value = newSeq[byte](valueLen)
|
|
if valueLen > 0:
|
|
copyMem(addr value[0], unsafeAddr mf.regions[0].data[pos], valueLen)
|
|
pos += valueLen
|
|
|
|
let timestamp = mf.readUint64(pos)
|
|
pos += 8
|
|
let deleted = mf.readByte(pos) != 0
|
|
|
|
return (true, Entry(key: key, value: value, timestamp: timestamp, deleted: deleted))
|
|
|
|
proc close*(sst: var SSTable) =
|
|
sst.mmapFile.close()
|
|
|
|
# ----------------------------------------------------------------------
|
|
# LSMTree API
|
|
# ----------------------------------------------------------------------
|
|
|
|
proc flushUnsafe(db: LSMTree) {.gcsafe.}
|
|
|
|
proc newLSMTree*(dir: string, memMaxSize: int = DefaultMemTableSize): LSMTree =
|
|
createDir(dir)
|
|
createDir(dir / "sstables")
|
|
|
|
var sstables: seq[SSTable] = @[]
|
|
var nextId = 1
|
|
|
|
# Load existing SSTables
|
|
for kind, path in walkDir(dir / "sstables"):
|
|
if kind == pcFile and path.endsWith(".sst"):
|
|
try:
|
|
var sst = loadSSTable(path)
|
|
let name = splitFile(path).name
|
|
sst.id = parseInt(name)
|
|
sstables.add(sst)
|
|
nextId = max(nextId, sst.id + 1)
|
|
except:
|
|
discard # skip corrupt SSTables
|
|
|
|
sstables.sort(proc(a, b: SSTable): int = cmp(b.id, a.id))
|
|
|
|
new(result)
|
|
initLock(result.lock)
|
|
initLock(result.walLock)
|
|
result.dir = dir
|
|
result.memTable = newMemTable(memMaxSize)
|
|
result.immutableMem = newMemTable(0)
|
|
result.sstables = sstables
|
|
result.wal = newWriteAheadLog(dir / "wal")
|
|
result.memMaxSize = memMaxSize
|
|
result.currentSeq = 0
|
|
result.nextSSTableId = nextId
|
|
|
|
# WAL crash recovery — replay unflushed entries into memTable
|
|
let walPath = dir / "wal" / "wal.log"
|
|
if fileExists(walPath):
|
|
var stream: FileStream = nil
|
|
try:
|
|
stream = newFileStream(walPath, fmRead)
|
|
if stream != nil:
|
|
var magic: uint32 = 0
|
|
var version: uint32 = 0
|
|
if stream.readData(addr magic, 4) == 4 and magic == WALMagic:
|
|
if stream.readData(addr version, 4) == 4:
|
|
while not stream.atEnd():
|
|
var kind: uint8 = 0
|
|
var timestamp: uint64 = 0
|
|
var keyLen: uint32 = 0
|
|
var valLen: uint32 = 0
|
|
if stream.readData(addr kind, 1) != 1: break
|
|
if stream.readData(addr timestamp, 8) != 8: break
|
|
if stream.readData(addr keyLen, 4) != 4: break
|
|
var key = newString(keyLen.int)
|
|
if keyLen > 0:
|
|
if stream.readData(addr key[0], keyLen.int) != keyLen.int: break
|
|
if stream.readData(addr valLen, 4) != 4: break
|
|
var value = newSeq[byte](valLen.int)
|
|
if valLen > 0:
|
|
if stream.readData(addr value[0], valLen.int) != valLen.int: break
|
|
case WalEntryKind(kind)
|
|
of wekPut:
|
|
if not result.memTable.put(key, value, timestamp):
|
|
result.flushUnsafe()
|
|
discard result.memTable.put(key, value, timestamp)
|
|
of wekDelete:
|
|
if not result.memTable.put(key, @[], timestamp, deleted = true):
|
|
result.flushUnsafe()
|
|
discard result.memTable.put(key, @[], timestamp, deleted = true)
|
|
of wekCommit:
|
|
discard
|
|
of wekCheckpoint:
|
|
discard
|
|
finally:
|
|
if stream != nil:
|
|
stream.close()
|
|
|
|
proc put*(db: LSMTree, key: string, value: seq[byte]) =
|
|
let ts = uint64(getMonoTime().ticks())
|
|
acquire(db.walLock)
|
|
db.wal.writePut(cast[seq[byte]](key), value, ts)
|
|
release(db.walLock)
|
|
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
if not db.memTable.put(key, value, ts):
|
|
if db.immutableMem.len > 0:
|
|
db.flushUnsafe()
|
|
db.immutableMem = db.memTable
|
|
db.memTable = newMemTable(db.memMaxSize)
|
|
discard db.memTable.put(key, value, ts)
|
|
|
|
proc delete*(db: LSMTree, key: string) =
|
|
let ts = uint64(getMonoTime().ticks())
|
|
acquire(db.walLock)
|
|
db.wal.writeDelete(cast[seq[byte]](key), ts)
|
|
release(db.walLock)
|
|
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
if not db.memTable.put(key, @[], ts, deleted = true):
|
|
if db.immutableMem.len > 0:
|
|
db.flushUnsafe()
|
|
db.immutableMem = db.memTable
|
|
db.memTable = newMemTable(db.memMaxSize)
|
|
discard db.memTable.put(key, @[], ts, deleted = true)
|
|
|
|
proc getUnsafe(db: LSMTree, key: string): (bool, seq[byte]) =
|
|
let (found, entry) = db.memTable.get(key)
|
|
if found:
|
|
if entry.deleted:
|
|
return (false, @[])
|
|
return (true, entry.value)
|
|
|
|
let (found2, entry2) = db.immutableMem.get(key)
|
|
if found2:
|
|
if entry2.deleted:
|
|
return (false, @[])
|
|
return (true, entry2.value)
|
|
|
|
# Search SSTables from newest to oldest
|
|
for i in countdown(db.sstables.high, db.sstables.low):
|
|
let sst = db.sstables[i]
|
|
if key < sst.minKey or key > sst.maxKey:
|
|
continue
|
|
if not sst.bloom.contains(cast[seq[byte]](key)):
|
|
continue
|
|
let (found3, entry3) = readSSTableEntry(sst, key)
|
|
if found3:
|
|
if entry3.deleted:
|
|
return (false, @[])
|
|
return (true, entry3.value)
|
|
|
|
return (false, @[])
|
|
|
|
proc get*(db: LSMTree, key: string): (bool, seq[byte]) =
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
return getUnsafe(db, key)
|
|
|
|
proc contains*(db: LSMTree, key: string): bool =
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
let (found, _) = getUnsafe(db, key)
|
|
return found
|
|
|
|
proc flushUnsafe(db: LSMTree) =
|
|
if db.immutableMem.len == 0 and db.memTable.len == 0:
|
|
return
|
|
|
|
# Flush immutable memtable if present, otherwise flush current memtable
|
|
var toFlush = db.immutableMem
|
|
if toFlush.len == 0:
|
|
toFlush = db.memTable
|
|
db.memTable = newMemTable(db.memMaxSize)
|
|
else:
|
|
db.immutableMem = newMemTable(0)
|
|
|
|
if toFlush.len == 0:
|
|
return
|
|
|
|
let path = db.dir / "sstables" / ($db.nextSSTableId & ".sst")
|
|
inc db.nextSSTableId
|
|
|
|
var sst = writeSSTable(toFlush.entries, path, level = 0)
|
|
sst.id = db.nextSSTableId - 1
|
|
db.sstables.add(sst)
|
|
# SSTables are kept in insertion order (newest last) so getUnsafe can search newest-first
|
|
|
|
db.wal.writeCommit(uint64(getMonoTime().ticks()))
|
|
db.wal.sync()
|
|
|
|
proc flush*(db: LSMTree) =
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
flushUnsafe(db)
|
|
|
|
proc close*(db: LSMTree) =
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
flushUnsafe(db)
|
|
for sst in db.sstables.mitems:
|
|
sst.close()
|
|
db.wal.close()
|
|
|
|
proc memTableSize*(db: LSMTree): int =
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
return db.memTable.len
|
|
|
|
proc sstableCount*(db: LSMTree): int =
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
return db.sstables.len
|
|
|
|
proc dir*(db: LSMTree): string =
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
return db.dir
|
|
|
|
proc scanMemTable*(db: LSMTree): seq[Entry] =
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
## Return all entries from memory (memTable + immutableMem)
|
|
result = @[]
|
|
for e in db.memTable.entries:
|
|
result.add(e)
|
|
for e in db.immutableMem.entries:
|
|
result.add(e)
|
|
|
|
proc scanAll*(db: LSMTree): seq[(string, seq[byte])] =
|
|
## Scan all active (non-deleted) entries from memory and SSTables.
|
|
## Used for shard data migration.
|
|
acquire(db.lock)
|
|
defer: release(db.lock)
|
|
|
|
var seen = initTable[string, bool]()
|
|
|
|
# Scan memtable first (most recent)
|
|
for e in db.memTable.entries:
|
|
if e.key notin seen:
|
|
seen[e.key] = true
|
|
if not e.deleted:
|
|
result.add((e.key, e.value))
|
|
|
|
# Scan immutable memtable
|
|
for e in db.immutableMem.entries:
|
|
if e.key notin seen:
|
|
seen[e.key] = true
|
|
if not e.deleted:
|
|
result.add((e.key, e.value))
|
|
|
|
# Scan SSTables from newest to oldest
|
|
for sst in db.sstables:
|
|
for key, offset in sst.index:
|
|
if key notin seen:
|
|
seen[key] = true
|
|
let (found, entry) = readSSTableEntry(sst, key)
|
|
if found and not entry.deleted:
|
|
result.add((entry.key, entry.value))
|