[
  {
    "path": "LICENSE",
    "content": "MIT License\n\nCopyright (c) 2025 Johnsoy.zhao\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "README.md",
    "content": "# 🚀 SmartStashDB: A High-Performance Key-Value Store\n\nWelcome to **SmartStashDB**, a blazing-fast, Go-powered key-value store built from scratch using **LSM-Tree**, **Skip-List**, and **Write-Ahead Logging (WAL)**. Designed for high throughput and low latency, SmartStashDB is perfect for applications demanding scalable, reliable, and efficient data storage.\n\n---\n\n## 🌟 Features\n\n- **High Performance**: Optimized for low-latency reads and writes, leveraging LSM-Tree and Skip-List for efficient data organization.\n- **Durability**: Write-Ahead Logging ensures no data is lost, even in the face of crashes.\n- **Scalability**: LSM-Tree architecture supports massive datasets with seamless compaction.\n- **Memory Efficiency**: Skip-List provides fast in-memory indexing with minimal overhead.\n- **Simple API**: Intuitive key-value operations for easy integration.\n- **Go-Powered**: Written in Go for concurrency, simplicity, and cross-platform support.\n\n---\n\n## 🛠️ Architecture\n\nSmartStashDB combines cutting-edge data structures and techniques to deliver top-tier performance:\n\n- **LSM-Tree**: Log-Structured Merge-Tree for write-heavy workloads, with background compaction to keep reads fast.\n- **Skip-List**: Probabilistic data structure for in-memory indexing, enabling O(log n) lookups.\n- **WAL**: Write-Ahead Logging for crash recovery and data durability.\n- **Compaction**: Periodic merging of SSTables to optimize storage and query performance.\n\n```\n[Client] --> [API: Get/Put/Delete] --> [MemTable (Skip-List)]\n                                             |\n                                             v\n                                        [WAL (Disk)]\n                                             |\n                                             v\n                                      [SSTables (LSM-Tree)]\n```\n\n---\n\n## 🚀 Getting Started\n\n### Prerequisites\n- **Go**: Version 1.18 or higher\n- A passion for high-performance systems! 😎\n\n### Installation\n1. Clone the repository:\n   ```bash\n   git clone https://github.com/johnsoy/SmartStashDB.git\n   cd SmartStashDB\n   ```\n2. Install dependencies:\n   ```bash\n   go mod tidy\n   ```\n3. Build and run:\n   ```bash\n   go build\n   ./SmartStashDB\n   ```\n\n### Example Usage\n```go\npackage main\n\nimport (\n    \"fmt\"\n    \"github.com/johnsoy/SmartStashDB\"\n)\n\nfunc main() {\n    // Initialize SmartStashDB\n    kv, err := SmartStashDB.NewSmartStashDB(\"./data\")\n    if err != nil {\n        panic(err)\n    }\n    defer kv.Close()\n\n    // Put key-value pair\n    kv.Put([]byte(\"key1\"), []byte(\"value1\"))\n\n    // Get value\n    value, err := kv.Get([]byte(\"key1\"))\n    if err != nil {\n        panic(err)\n    }\n    fmt.Printf(\"Key: key1, Value: %s\\n\", value)\n\n    // Delete key\n    kv.Delete([]byte(\"key1\"))\n}\n```\n\n---\n\n## 📊 Performance\n\nSmartStashDB is designed for speed and scalability. Preliminary benchmarks (on a standard laptop with SSD):\n- **Write Throughput**: ~500,000 ops/sec\n- **Read Throughput**: ~600,000 ops/sec\n- **Latency**: < 1ms for 99th percentile reads/writes\n\nRun benchmarks yourself:\n```bash\ngo test -bench=.\n```\n\n---\n\n## 🛠️ Configuration\n\nCustomize SmartStashDB via the `config.yaml` file:\n```yaml\ndata_dir: \"./data\"          # Storage directory\nmemtable_size: 1048576      # Max MemTable size (bytes)\ncompaction_interval: 60     # Compaction interval (seconds)\nwal_flush_interval: 10      # WAL flush interval (seconds)\n```\n\nLoad config programmatically:\n```go\nkv, err := SmartStashDB.NewSmartStashDBWithConfig(\"config.yaml\")\n```\n\n---\n\n## 🤝 Contributing\n\nContributions are welcome! Whether it's bug fixes, new features, or documentation improvements, here's how to get started:\n1. Fork the repository.\n2. Create a feature branch: `git checkout -b feature/awesome-feature`.\n3. Commit your changes: `git commit -m \"Add awesome feature\"`.\n4. Push to the branch: `git push origin feature/awesome-feature`.\n5. Open a Pull Request.\n\nPlease read our [CONTRIBUTING.md](CONTRIBUTING.md) for more details.\n\n---\n\n## 📜 License\n\nSmartStashDB is licensed under the [MIT License](LICENSE). Feel free to use, modify, and distribute it as you see fit!\n\n---\n\n## 📫 Contact\n\n- **GitHub**: [Johnsoy](https://github.com/Johonsoy)\n- **Email**: [15520754767@163.com]\n\nStar ⭐ this repo if you find SmartStashDB awesome, and let's build the fastest KV store together! 🚀\n"
  },
  {
    "path": "const/Reader.go",
    "content": "package _const\n\nimport (\n\t\"SmartStashDB/storage\"\n\t\"io\"\n)\n\ntype Reader struct {\n\tAllSegmentReader []*storage.SegmentReader\n\tProgress         int\n}\n\nfunc (r *Reader) Next() ([]byte, *storage.ChunkPosition, error) {\n\tif r.Progress >= len(r.AllSegmentReader) {\n\t\treturn nil, nil, io.EOF\n\t}\n\tdata, chunkPos, err := r.AllSegmentReader[r.Progress].Next()\n\tif err == io.EOF {\n\t\tr.Progress++\n\t\treturn r.Next()\n\t}\n\treturn data, chunkPos, err\n}\n"
  },
  {
    "path": "const/constant.go",
    "content": "package _const\n\nconst (\n\tB  = 1\n\tKB = 1024 * B\n\tMB = 1024 * KB\n\tGB = 1024 * MB\n)\n\nconst (\n\t// 单个Block 32KB\n\tBlockSize = 32 * KB\n)\n\nconst (\n\tChunkHeadSize = 7\n)\n\nconst (\n\tFirstSegmentFileId  = 1\n\tsegmentFileModePerm = 0644\n)\n\nfunc ExecDir() string {\n\n\treturn \"\"\n}\n"
  },
  {
    "path": "const/error.go",
    "content": "package _const\n\nimport (\n\t\"errors\"\n)\n\nvar (\n\tErrDatabaseIsUsing       = errors.New(\"the database directory is used by another process\")\n\tErrorDBClosed            = errors.New(\"the database is closed\")\n\tErrorReadOnlyBatch       = errors.New(\"the read-only batch exists\")\n\tErrorBatchCommited       = errors.New(\"the batch commited\")\n\tErrorKeyNotFound         = errors.New(\"key not found\")\n\tErrorKeyIsEmpty          = errors.New(\"the key is empty\")\n\tErrorFileExtError        = errors.New(\"segmentFileExt must not start with '.'\")\n\tErrorDataToLarge         = errors.New(\"data is too large\")\n\tErrorPendingSizeTooLarge = errors.New(\"pending size is too large\")\n\tErrClosed                = errors.New(\"closed\")\n)\n"
  },
  {
    "path": "go.mod",
    "content": "module SmartStashDB\n\ngo 1.22\n\nrequire (\n\tgithub.com/bwmarrin/snowflake v0.3.0\n\tgithub.com/dgraph-io/badger v1.6.0\n\tgithub.com/gofrs/flock v0.8.1\n\tgithub.com/hashicorp/golang-lru/v2 v2.0.7\n)\n\nrequire (\n\tgithub.com/pkg/errors v0.9.1 // indirect\n\tgithub.com/stretchr/testify v1.10.0 // indirect\n\tgolang.org/x/net v0.7.0 // indirect\n\tgolang.org/x/sys v0.5.0 // indirect\n\tgopkg.in/yaml.v2 v2.2.8 // indirect\n)\n"
  },
  {
    "path": "main.go",
    "content": "package main\n\nimport (\n\t_const \"SmartStashDB/const\"\n\t\"SmartStashDB/storage\"\n)\n\nfunc main() {\n\toptions := storage.DefaultOptions\n\n\toptions.DirPath = _const.ExecDir() + \"/data\"\n\n\tdb, err := storage.OpenDB(options)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\tdefer func() {\n\t\t_ = db.Close()\n\t}()\n\tkey := \"adasdsa\"\n\tvalue := \"asdbsadsd\"\n\terr = db.Put(key, value, nil)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\n\tnewValue, err := db.Get(key)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\tprint(string(newValue))\n}\n"
  },
  {
    "path": "storage/SegmentReader.go",
    "content": "package storage\n\nimport (\n\t_const \"SmartStashDB/const\"\n\t\"io\"\n)\n\ntype SegmentReader struct {\n\tseg         *SegmentFile\n\tblockidx    uint32\n\tchunkoffset uint32\n}\n\nfunc (s *SegmentReader) Next() ([]byte, *ChunkPosition, error) {\n\tif s.seg.closed {\n\t\treturn nil, nil, io.EOF\n\t}\n\n\tcurChunk := &ChunkPosition{\n\t\tSegmentFileId: s.seg.segmentFileId,\n\t\tBlockIndex:    s.blockidx,\n\t\tChunkOffset:   s.chunkoffset,\n\t}\n\tdata, nextChunk, err := s.seg.readInternal(curChunk.BlockIndex, curChunk.ChunkOffset)\n\tif err != nil {\n\t\treturn nil, nil, err\n\t}\n\tcurChunk.ChunkSize = nextChunk.BlockIndex*_const.BlockSize + nextChunk.ChunkOffset -\n\t\t(s.chunkoffset*_const.BlockSize + curChunk.ChunkOffset)\n\ts.blockidx = nextChunk.BlockIndex\n\ts.chunkoffset = curChunk.ChunkOffset\n\treturn data, curChunk, nil\n}\n"
  },
  {
    "path": "storage/batch.go",
    "content": "package storage\n\nimport (\n\t_const \"SmartStashDB/const\"\n\t\"sync\"\n)\nimport \"github.com/bwmarrin/snowflake\"\n\nfunc makeBatch() interface{} {\n\tnode, err := snowflake.NewNode(1)\n\tif err != nil {\n\t\tpanic(err)\n\t}\n\treturn &Batch{\n\t\toptions: DefaultBatchOptions,\n\t\tm:       sync.RWMutex{},\n\t\tbatchId: node,\n\t}\n}\n\ntype Batch struct {\n\tdb            *DB\n\tpendingWrites map[string]*LogRecord\n\toptions       BatchOptions\n\tm             sync.RWMutex\n\tcommited      bool\n\tbatchId       *snowflake.Node\n}\n\nfunc (batch *Batch) reset() {\n\n}\n\nfunc (batch *Batch) init(readOnly bool, sync bool, db *DB) *Batch {\n\tbatch.db = db\n\tbatch.options.ReadOnly = readOnly\n\tbatch.options.Sync = sync\n\tbatch.lock()\n\treturn batch\n}\n\nfunc (batch *Batch) lock() {\n\tif batch.options.ReadOnly {\n\t\tbatch.db.m.RLock()\n\t} else {\n\t\tbatch.db.m.Lock()\n\t}\n}\n\nfunc (batch *Batch) writePendingWrites() *Batch {\n\tbatch.pendingWrites = make(map[string]*LogRecord)\n\treturn batch\n}\n\nfunc (batch *Batch) put(key []byte, value []byte) error {\n\tif len(key) == 0 {\n\t\treturn _const.ErrorKeyIsEmpty\n\t}\n\n\tif batch.db.Closed {\n\t\treturn _const.ErrorDBClosed\n\t}\n\n\tif batch.options.ReadOnly {\n\t\treturn _const.ErrorReadOnlyBatch\n\t}\n\tbatch.m.Lock()\n\tdefer batch.m.Unlock()\n\tbatch.pendingWrites[string(key)] = &LogRecord{\n\t\tKey:   key,\n\t\tValue: value,\n\t\tType:  LogRecordNormal,\n\t}\n\treturn nil\n}\n\nfunc (batch *Batch) unLock() {\n\tif batch.options.ReadOnly {\n\t\tbatch.db.m.RUnlock()\n\t} else {\n\t\tbatch.db.m.Unlock()\n\t}\n\n}\n\nfunc (batch *Batch) commit(w *WriteOptions) error {\n\tif w == nil {\n\t\tw = &WriteOptions{\n\t\t\tSync:       false,\n\t\t\tDisableWal: false,\n\t\t}\n\t}\n\tdefer batch.unLock()\n\tif batch.db.Closed {\n\t\treturn _const.ErrorDBClosed\n\t}\n\n\tif batch.options.ReadOnly || len(batch.pendingWrites) == 0 {\n\t\treturn nil\n\t}\n\n\tbatch.m.Lock()\n\tdefer batch.m.Unlock()\n\tif batch.commited {\n\t\treturn _const.ErrorBatchCommited\n\t}\n\n\tif err := batch.db.waitMemTableSpace(); err != nil {\n\t\treturn err\n\t}\n\n\tbatchId := batch.batchId.Generate()\n\tif err := batch.db.activeMem.putBatch(batch.pendingWrites, batchId, w); err != nil {\n\t\treturn err\n\t}\n\tbatch.commited = true\n\treturn nil\n}\n\nfunc (batch *Batch) Get(key []byte) ([]byte, error) {\n\tif len(key) == 0 {\n\t\treturn nil, _const.ErrorKeyIsEmpty\n\t}\n\n\tif batch.db.Closed {\n\t\treturn nil, _const.ErrorDBClosed\n\t}\n\n\tif batch.pendingWrites != nil {\n\t\tbatch.m.RLock()\n\t\tdefer batch.m.RUnlock()\n\t\tif record := batch.pendingWrites[string(key)]; record != nil {\n\t\t\tif record.Type == LogRecordDeleted {\n\t\t\t\treturn nil, _const.ErrorKeyNotFound\n\t\t\t}\n\t\t\treturn record.Value, nil\n\t\t}\n\t}\n\n\ttables := batch.db.getMemTables()\n\n\tfor _, table := range tables {\n\t\tdeleted, value := table.get(key)\n\t\tif deleted {\n\t\t\treturn nil, _const.ErrorKeyNotFound\n\t\t}\n\t\tif len(value) != 0 {\n\t\t\treturn value, nil\n\t\t}\n\t}\n\n\treturn nil, _const.ErrorKeyNotFound\n}\n\nfunc (batch *Batch) delete(key []byte) error {\n\tif len(key) == 0 {\n\t\treturn _const.ErrorKeyIsEmpty\n\t}\n\n\tif batch.db.Closed {\n\t\treturn _const.ErrorDBClosed\n\t}\n\n\tif batch.options.ReadOnly {\n\t\treturn _const.ErrorReadOnlyBatch\n\t}\n\tbatch.m.Lock()\n\tbatch.pendingWrites[string(key)] = &LogRecord{\n\t\tKey:  key,\n\t\tType: LogRecordDeleted,\n\t}\n\tbatch.m.Unlock()\n\treturn nil\n}\n"
  },
  {
    "path": "storage/chunk.go",
    "content": "package storage\n\ntype ChunkType = byte\n\nconst (\n\tChunkTypeFull ChunkType = iota\n\tChunkTypeStart\n\tChunkTypeMiddle\n\tChunkTypeEnd\n)\n\ntype ChunkPosition struct {\n\tSegmentFileId SegmentFileId\n\tBlockIndex    uint32\n\tChunkOffset   uint32\n\tChunkSize     uint32\n}\n"
  },
  {
    "path": "storage/db.go",
    "content": "package storage\n\nimport (\n\t_const \"SmartStashDB/const\"\n\t\"errors\"\n\t\"github.com/gofrs/flock\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"sync\"\n)\n\nconst (\n\tFileLockName = \"FLOCK\"\n)\n\ntype DB struct {\n\tm            sync.RWMutex\n\tactiveMem    *MemTable   // Active memory\n\timmutableMem []*MemTable // Immutable memory\n\tClosed       bool\n\tbatchPool    sync.Pool\n}\n\nfunc (db *DB) Close() error {\n\tdb.m.Lock()\n\tdefer db.m.Unlock()\n\n\tfor _, table := range db.immutableMem {\n\t\terr := table.close()\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\tif err := db.activeMem.close(); err != nil {\n\t\treturn err\n\t}\n\tdb.Closed = true\n\treturn nil\n}\n\nfunc (db *DB) Put(key string, value string, options *WriteOptions) error {\n\tbatch := db.batchPool.Get().(*Batch)\n\tdefer func() {\n\t\tbatch.reset()\n\t\tdb.batchPool.Put(batch)\n\t}()\n\tbatch.init(false, false, db).writePendingWrites()\n\terr := batch.put([]byte(key), []byte(value))\n\tif err != nil {\n\t\tbatch.unLock()\n\t\treturn err\n\t}\n\treturn batch.commit(options)\n}\n\nfunc (db *DB) waitMemTableSpace() error {\n\tif db.activeMem.isFull() {\n\t\treturn nil\n\t}\n\tdb.immutableMem = append(db.immutableMem, db.activeMem)\n\toption := db.activeMem.option\n\toption.id++\n\ttable, err := openMemTable(option)\n\tif err != nil {\n\t\treturn err\n\t}\n\tdb.activeMem = table\n\treturn nil\n}\n\nfunc (db *DB) Get(key string) ([]byte, error) {\n\tbatch := db.batchPool.Get().(*Batch)\n\tbatch.init(true, false, db)\n\tdefer func() {\n\t\t_ = batch.commit(nil)\n\t\tbatch.reset()\n\t\tdb.batchPool.Put(batch)\n\t}()\n\treturn batch.Get([]byte(key))\n}\n\nfunc (db *DB) getMemTables() []*MemTable {\n\treturn db.immutableMem\n}\n\nfunc (db *DB) Delete(key []byte, options *WriteOptions) error {\n\tif len(key) == 0 {\n\t\treturn _const.ErrorKeyIsEmpty\n\t}\n\tbatch := db.batchPool.Get().(*Batch)\n\tbatch.init(false, false, db)\n\tdefer func() {\n\t\tbatch.reset()\n\t\tdb.batchPool.Put(batch)\n\t}()\n\tif err := batch.delete(key); err != nil {\n\t\tbatch.unLock()\n\t\treturn err\n\t}\n\treturn batch.commit(options)\n}\n\nfunc OpenDB(options Options) (*DB, error) {\n\n\t// Check if file existed.\n\tif _, err := os.Stat(options.DirPath); err != nil {\n\t\tif err := os.Mkdir(options.DirPath, os.ModePerm); err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\tlock, err := flock.New(filepath.Join(options.DirPath, FileLockName)).TryLock()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif !lock {\n\t\treturn nil, errors.New(\"file locked\")\n\t}\n\n\tmemTables, err := openAllMemTables(options)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tdb := &DB{\n\t\tactiveMem:    memTables[len(memTables)-1],\n\t\timmutableMem: memTables,\n\t\tbatchPool:    sync.Pool{New: makeBatch},\n\t}\n\treturn db, nil\n}\n"
  },
  {
    "path": "storage/logrecord.go",
    "content": "package storage\n\nimport (\n\t\"encoding/binary\"\n)\n\ntype LogRecordType = byte\n\nconst (\n\tLogRecordNormal = iota\n\tLogRecordDeleted\n\tLogRecordBatchEnd\n\tMaxLogRecordLength = 1 + binary.MaxVarintLen64*2\n)\n\ntype LogRecord struct {\n\tKey     []byte\n\tValue   []byte\n\tType    LogRecordType\n\tBatchId uint64\n}\n\nfunc NewLogRecord() *LogRecord {\n\treturn &LogRecord{}\n}\n\n// Encode Serialize LogRecord, header + batchId + keySize + valueSize + key + value /*\nfunc (logRecord *LogRecord) Encode() []byte {\n\theader := make([]byte, MaxLogRecordLength)\n\theader[0] = logRecord.Type\n\n\tindex := 1\n\n\tindex += binary.PutUvarint(header[index:], logRecord.BatchId)\n\tindex += binary.PutVarint(header[index:], int64(len(logRecord.Key)))\n\tindex += binary.PutVarint(header[index:], int64(len(logRecord.Value)))\n\n\tvalue := make([]byte, index+len(logRecord.Key)+len(logRecord.Value))\n\n\t// copy header.\n\tcopy(value, header[:index])\n\n\tcopy(value[index:], logRecord.Key)\n\n\tcopy(value[index+len(logRecord.Key):], logRecord.Value)\n\n\treturn value\n}\n\nfunc (logRecord *LogRecord) Decode(b []byte) {\n\tlogRecord.Type = b[0]\n\n\tindex := 1\n\tn := 0\n\tlogRecord.BatchId, n = binary.Uvarint(b[index:])\n\tindex += n\n\tkeyLength, n := binary.Uvarint(b[index:])\n\tindex += n\n\n\tvalueLength, n := binary.Uvarint(b[index:])\n\tindex += n\n\n\tkey := make([]byte, keyLength)\n\n\tvalue := make([]byte, valueLength)\n\n\tcopy(key, b[index:index+int(keyLength)])\n\tindex += int(keyLength)\n\n\tcopy(value, b[index:index+int(valueLength)])\n\n\tlogRecord.Key = key\n\tlogRecord.Value = value\n\n}\n"
  },
  {
    "path": "storage/memtable.go",
    "content": "package storage\n\nimport (\n\t\"fmt\"\n\t\"github.com/bwmarrin/snowflake\"\n\t\"github.com/dgraph-io/badger/skl\"\n\t\"github.com/dgraph-io/badger/y\"\n\t\"io\"\n\t\"math\"\n\t\"os\"\n\t\"sort\"\n\t\"sync\"\n)\n\nconst (\n\tinitTableId = 1\n\n\twalFileExt = \".MEM.%d\"\n)\n\ntype MemTable struct {\n\toption memTableOptions\n\n\tmu sync.RWMutex\n\n\tskl *skl.Skiplist\n\n\ttinyWal *TinyWAL\n}\n\ntype memTableOptions struct {\n\tsklMemSize      uint32 // skip-list memory size.\n\tid              int    // skip-list memory id.\n\twalDir          string // file dir.\n\twalCacheSize    uint32 // wal cache size.\n\twalIsSync       bool   // whether to flush the disk immediately.\n\twalBytesPerSync uint32 // how bytes to flush the disk.\n}\n\nfunc openAllMemTables(options WalOptions) ([]*MemTable, error) {\n\tdir, err := os.ReadDir(options.DirPath)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tvar tableIds []int\n\n\tfor _, file := range dir {\n\t\tif file.IsDir() {\n\t\t\tcontinue\n\t\t}\n\t\tvar id int\n\t\tvar prefix int\n\n\t\t_, err = fmt.Sscanf(file.Name(), \"memtable_%d\"+walFileExt, &prefix, &id)\n\t\tif err != nil {\n\t\t\tcontinue\n\t\t}\n\t\ttableIds = append(tableIds, id)\n\t}\n\n\tif len(tableIds) == 0 {\n\t\ttableIds = append(tableIds, initTableId)\n\t}\n\n\tsort.Ints(tableIds)\n\n\ttables := make([]*MemTable, len(tableIds))\n\n\tfor i, id := range tableIds {\n\t\ttable, err := openMemTable(memTableOptions{\n\t\t\tsklMemSize:      options.MemTableSize,\n\t\t\tid:              id,\n\t\t\twalDir:          options.DirPath,\n\t\t\twalIsSync:       options.Sync,\n\t\t\twalBytesPerSync: options.BytesPerSync,\n\t\t})\n\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\ttables[i] = table\n\t}\n\n\treturn tables, nil\n}\n\nfunc openMemTable(option memTableOptions) (*MemTable, error) {\n\tskipList := skl.NewSkiplist(int64(option.sklMemSize * 2))\n\n\ttable := &MemTable{\n\t\toption: option,\n\t\tskl:    skipList,\n\t}\n\twal, err := OpenTinyWAL(WalOptions{\n\t\tDirPath:        option.walDir,\n\t\tMemTableSize:   math.MaxInt32,\n\t\tsegmentFileExt: fmt.Sprintf(walFileExt, option.id),\n\t\tSync:           option.walIsSync,\n\t\tBytesPerSync:   option.walBytesPerSync,\n\t\tBlockCache:     option.walCacheSize,\n\t})\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\ttable.tinyWal = wal\n\n\tindexRecords := make(map[uint64][]*LogRecord)\n\n\treader := wal.NewReader()\n\n\tfor {\n\t\tdata, _, err := reader.Next()\n\t\tif err != nil {\n\t\t\tif err == io.EOF {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\treturn nil, err\n\t\t}\n\n\t\trecord := NewLogRecord()\n\t\trecord.Decode(data)\n\t\tif record.Type == LogRecordBatchEnd {\n\t\t\tbatchId, err := snowflake.ParseBytes(record.Key)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\n\t\t\tfor _, idxRecord := range indexRecords[uint64(batchId)] {\n\t\t\t\ttable.skl.Put(y.KeyWithTs(idxRecord.Key, 0),\n\t\t\t\t\ty.ValueStruct{\n\t\t\t\t\t\tMeta:  idxRecord.Type,\n\t\t\t\t\t\tValue: idxRecord.Value,\n\t\t\t\t\t})\n\t\t\t}\n\t\t\tdelete(indexRecords, uint64(batchId))\n\n\t\t} else {\n\t\t\tindexRecords[record.BatchId] = append(indexRecords[record.BatchId], record)\n\t\t}\n\t}\n\n\treturn table, nil\n}\n\nfunc (mt *MemTable) get(key []byte) (bool, []byte) {\n\tmt.mu.RLock()\n\tdefer mt.mu.RUnlock()\n\n\tvalueStruct := mt.skl.Get(y.KeyWithTs(key, 0))\n\tdeleted := valueStruct.Meta == LogRecordDeleted\n\n\treturn deleted, valueStruct.Value\n}\n\nfunc (mt *MemTable) isFull() bool {\n\treturn mt.skl.MemSize() >= int64(mt.option.sklMemSize)\n}\n\nfunc (mt *MemTable) putBatch(records map[string]*LogRecord, batchId snowflake.ID, options *WriteOptions) error {\n\tif options == nil || options.DisableWal {\n\t\tfor _, record := range records {\n\t\t\trecord.BatchId = uint64(batchId)\n\t\t\tif err := mt.tinyWal.PendingWrites(record.Encode()); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t}\n\t\trecord := NewLogRecord()\n\t\trecord.Key = batchId.Bytes()\n\t\trecord.Type = LogRecordBatchEnd\n\n\t\tif err := mt.tinyWal.PendingWrites(record.Encode()); err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tif err := mt.tinyWal.WriteAll(); err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tif options != nil && options.Sync && mt.option.walIsSync {\n\t\t\tif err := mt.tinyWal.Sync(); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t}\n\t}\n\n\tmt.mu.Lock()\n\tfor key, record := range records {\n\t\tmt.skl.Put(y.KeyWithTs([]byte(key), 0),\n\t\t\ty.ValueStruct{\n\t\t\t\tMeta:  record.Type,\n\t\t\t\tValue: record.Value,\n\t\t\t})\n\t}\n\tmt.mu.Unlock()\n\treturn nil\n}\n\nfunc (mt *MemTable) close() error {\n\tif mt.skl != nil {\n\t\treturn mt.tinyWal.close()\n\t}\n\treturn nil\n}\n"
  },
  {
    "path": "storage/options.go",
    "content": "package storage\n\nimport (\n\t_const \"SmartStashDB/const\"\n\t\"os\"\n)\n\ntype WalOptions struct {\n\tDirPath        string\n\tMemTableSize   uint64\n\tsegmentFileExt string\n\tSync           bool\n\tBytesPerSync   uint64\n\tBlockCache     uint32\n}\n\ntype BatchOptions struct {\n\tReadOnly bool\n\tSync     bool\n}\n\nvar DefaultOptions = WalOptions{\n\tDirPath:      tempDBDir(),\n\tMemTableSize: 64 * _const.MB,\n\tBlockCache:   0,\n\tSync:         false,\n\tBytesPerSync: 0,\n}\n\nvar DefaultBatchOptions = BatchOptions{\n\tReadOnly: false,\n\tSync:     true,\n}\n\nfunc tempDBDir() string {\n\ttemp, _ := os.MkdirTemp(\"\", \"db-temp\")\n\treturn temp\n}\n\ntype WriteOptions struct {\n\tSync       bool\n\tDisableWal bool\n}\n"
  },
  {
    "path": "storage/pool.go",
    "content": "package storage\n\nimport (\n\t\"bytes\"\n\t\"sync\"\n)\n\ntype bufferPool struct {\n\tbuffer sync.Pool\n}\n\nfunc (p *bufferPool) Get() *bytes.Buffer {\n\treturn p.buffer.Get().(*bytes.Buffer)\n}\n\nfunc (p *bufferPool) Put(buffer *bytes.Buffer) {\n\tp.buffer.Put(buffer)\n}\n\nvar DefaultBuffer = newBufferPool()\n\nfunc newBufferPool() *bufferPool {\n\treturn &bufferPool{\n\t\tbuffer: sync.Pool{\n\t\t\tNew: func() any { return new(bytes.Buffer) },\n\t\t},\n\t}\n}\n"
  },
  {
    "path": "storage/segmentfile.go",
    "content": "package storage\n\nimport (\n\t_const \"SmartStashDB/const\"\n\t\"bytes\"\n\t\"encoding/binary\"\n\t\"fmt\"\n\tlru \"github.com/hashicorp/golang-lru/v2\"\n\t\"hash/crc32\"\n\t\"os\"\n\t\"path/filepath\"\n)\n\ntype SegmentFileId = uint32\n\ntype SegmentFile struct {\n\tsegmentFileId SegmentFileId\n\n\tfd *os.File\n\n\tlastBlockIndex uint32\n\n\tlastBlockSize uint32\n\n\theader []byte\n\n\tclosed bool\n\n\tlocalCache *lru.Cache[uint32, []byte]\n}\n\nfunc (f *SegmentFile) readInternal(index uint32, offset uint32) ([]byte, *ChunkPosition, error) {\n\n\treturn nil, nil, nil\n}\n\nfunc (f *SegmentFile) NewSegmentReader() *SegmentReader {\n\treturn &SegmentReader{\n\t\tseg:         f,\n\t\tblockidx:    0,\n\t\tchunkoffset: 0,\n\t}\n}\n\nfunc (f *SegmentFile) Close() error {\n\tif f.closed {\n\t\treturn nil\n\t}\n\tf.closed = true\n\treturn f.fd.Close()\n}\n\nfunc (f *SegmentFile) Size() int64 {\n\treturn int64(f.lastBlockIndex*_const.BlockSize + f.lastBlockSize)\n}\n\nfunc (f *SegmentFile) Sync() error {\n\treturn f.fd.Sync()\n}\n\nfunc (f *SegmentFile) Write(data []byte) (*ChunkPosition, error) {\n\tif f.closed {\n\t\treturn nil, _const.ErrClosed\n\t}\n\tindex := f.lastBlockIndex\n\tsize := f.lastBlockSize\n\n\tvar err error\n\tbuffer := DefaultBuffer.Get()\n\tdefer func() {\n\t\tDefaultBuffer.Put(buffer)\n\t}()\n\twriteBuffer, err := f.writeBuffer(data, buffer)\n\tif err != nil {\n\t\tf.lastBlockIndex = index\n\t\tf.lastBlockSize = size\n\t\treturn nil, err\n\t}\n\terr = f.writeBuffer2File(buffer)\n\tif err != nil {\n\t\tf.lastBlockIndex = index\n\t\tf.lastBlockSize = size\n\t\treturn nil, err\n\t}\n\treturn writeBuffer, nil\n}\n\nfunc (f *SegmentFile) WriteAll(writes [][]byte) (position []*ChunkPosition, err error) {\n\tif f.closed {\n\t\treturn nil, _const.ErrClosed\n\t}\n\n\tindex := f.lastBlockIndex\n\tlastBlockSize := f.lastBlockSize\n\n\tbuffer := DefaultBuffer.Get()\n\n\tdefer func() {\n\t\tif err != nil {\n\t\t\tf.lastBlockIndex = index\n\t\t\tf.lastBlockSize = lastBlockSize\n\t\t}\n\t\tDefaultBuffer.Put(buffer)\n\t}()\n\n\tpositions := make([]*ChunkPosition, len(writes))\n\tfor i, data := range writes {\n\t\tpos, err := f.writeBuffer(data, buffer)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tpositions[i] = pos\n\t}\n\n\tif err := f.writeBuffer2File(buffer); err != nil {\n\t\treturn nil, err\n\t}\n\treturn positions, nil\n}\n\nfunc (f *SegmentFile) writeBuffer(bytes []byte, buffer *bytes.Buffer) (*ChunkPosition, error) {\n\tif f.closed {\n\t\treturn nil, _const.ErrClosed\n\t}\n\n\tpadding := uint32(0)\n\n\t// Pre-grow the buffer for better performance\n\ttotalWriteSize := len(bytes) + int(_const.ChunkHeadSize)*2 // Estimate needed size\n\tbuffer.Grow(totalWriteSize)\n\n\tif f.lastBlockSize+_const.ChunkHeadSize >= _const.BlockSize {\n\t\tsize := _const.BlockSize - f.lastBlockSize\n\t\t_, err := buffer.Write(make([]byte, size))\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tpadding += size\n\t\tf.lastBlockIndex++\n\t\tf.lastBlockSize = 0\n\t}\n\n\tposition := &ChunkPosition{\n\t\tSegmentFileId: f.segmentFileId,\n\t\tBlockIndex:    f.lastBlockIndex,\n\t\tChunkOffset:   f.lastBlockSize,\n\t}\n\n\tdataLen := uint32(len(bytes))\n\n\tif f.lastBlockSize+_const.ChunkHeadSize <= _const.BlockSize {\n\t\terr := f.appendChunk2Buffer(buffer, bytes, ChunkTypeFull)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tposition.ChunkSize = dataLen + _const.ChunkHeadSize\n\t} else {\n\t\t// Split data into many chunks across blocks\n\t\tvar (\n\t\t\tremainingDataSize        = dataLen\n\t\t\tcurBlockSize             = f.lastBlockSize\n\t\t\tchunkNum          uint32 = 0\n\t\t)\n\n\t\tfor remainingDataSize > 0 {\n\t\t\tchunkType := ChunkTypeMiddle\n\t\t\tif remainingDataSize == dataLen {\n\t\t\t\tchunkType = ChunkTypeStart\n\t\t\t}\n\t\t\tfreeSize := _const.BlockSize - curBlockSize - _const.ChunkHeadSize\n\t\t\tif freeSize >= remainingDataSize {\n\t\t\t\tfreeSize = remainingDataSize\n\t\t\t\tchunkType = ChunkTypeEnd\n\t\t\t}\n\t\t\terr := f.appendChunk2Buffer(buffer, bytes[dataLen-remainingDataSize:dataLen-remainingDataSize+freeSize], chunkType)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tchunkNum++\n\t\t\tremainingDataSize -= freeSize\n\t\t\tcurBlockSize = (curBlockSize + _const.ChunkHeadSize + freeSize) % _const.BlockSize\n\t\t}\n\t\tposition.ChunkSize = chunkNum*_const.ChunkHeadSize + dataLen\n\t}\n\n\treturn position, nil\n}\n\nfunc (f *SegmentFile) writeBuffer2File(buffer *bytes.Buffer) error {\n\tif f.lastBlockSize > _const.BlockSize {\n\t\tpanic(\"lastBlockSize exceeded BlockSize\")\n\t}\n\t_, err := f.fd.Write(buffer.Bytes())\n\treturn err\n}\n\nfunc (f *SegmentFile) appendChunk2Buffer(buffer *bytes.Buffer, data []byte, cType ChunkType) error {\n\t// 设置header中的长度\n\tbinary.LittleEndian.PutUint16(f.header[4:6], uint16(len(data)))\n\t// 设置header中的类型\n\tf.header[6] = cType\n\n\t// 对 len + type + data 求 checksum\n\tsum := crc32.ChecksumIEEE(f.header[4:])\n\tsum = crc32.Update(sum, crc32.IEEETable, data)\n\t// 设置header中的校验和\n\tbinary.LittleEndian.PutUint32(f.header[:4], sum)\n\t//将一个完整的chunk写入buf中（header + payload 就是一个chunk）\n\t_, err := buffer.Write(f.header[:])\n\tif err != nil {\n\t\treturn err\n\t}\n\t_, err = buffer.Write(data)\n\tif err != nil {\n\t\treturn err\n\t}\n\treturn nil\n}\n\nfunc segmentFileName(dir, ext string, id uint32) string {\n\treturn filepath.Join(dir, fmt.Sprintf(\"%010d\"+ext, id))\n}\n\nfunc openSegmentFile(dir string, ext string, id uint32, localCache *lru.Cache[uint32, []byte]) (*SegmentFile, error) {\n\tpath := segmentFileName(dir, ext, id)\n\tfd, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0666)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tdefer func() {\n\t\tif err != nil {\n\t\t\t_ = fd.Close()\n\t\t}\n\t}()\n\n\tstat, err := fd.Stat()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tsize := stat.Size()\n\treturn &SegmentFile{\n\t\tsegmentFileId:  id,\n\t\tfd:             fd,\n\t\tlastBlockIndex: uint32(size / _const.BlockSize),\n\t\tlastBlockSize:  uint32(size % _const.BlockSize),\n\t\theader:         make([]byte, _const.ChunkHeadSize),\n\t\tlocalCache:     localCache,\n\t}, nil\n}\n"
  },
  {
    "path": "storage/tinywal.go",
    "content": "package storage\n\nimport (\n\t_const \"SmartStashDB/const\"\n\t\"errors\"\n\t\"fmt\"\n\tlru \"github.com/hashicorp/golang-lru/v2\"\n\t\"io/fs\"\n\t\"os\"\n\t\"sort\"\n\t\"strings\"\n\t\"sync\"\n)\n\ntype TinyWAL struct {\n\toption           WalOptions\n\tmutex            sync.RWMutex\n\tactiveSegment    *SegmentFile\n\timmutableSegment map[SegmentFileId]*SegmentFile\n\tlocalCache       *lru.Cache[uint32, []byte]\n\tbyteWrite        uint64\n\n\tpendingWritesLock sync.Mutex\n\tpendingWrites     [][]byte\n\tpendingWritesSize uint64\n}\n\nfunc (w *TinyWAL) close() error {\n\tw.mutex.Lock()\n\tdefer w.mutex.Unlock()\n\n\tif w.localCache != nil {\n\t\tw.localCache.Purge()\n\t}\n\n\tfor _, segment := range w.immutableSegment {\n\t\tif segment != nil {\n\t\t\tif err := segment.Close(); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t}\n\t}\n\tw.immutableSegment = nil\n\treturn w.activeSegment.Close()\n}\n\nfunc OpenTinyWAL(option WalOptions) (*TinyWAL, error) {\n\tif strings.HasPrefix(option.segmentFileExt, \".\") {\n\t\treturn nil, errors.New(option.segmentFileExt + \" is not allowed\")\n\t}\n\n\terr := os.MkdirAll(option.DirPath, fs.ModePerm)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\ttinyWAL := &TinyWAL{\n\t\toption:           option,\n\t\timmutableSegment: make(map[SegmentFileId]*SegmentFile),\n\t\tactiveSegment:    nil,\n\t}\n\n\tif option.BlockCache > 0 {\n\t\tblockNum := option.BlockCache / _const.BlockSize\n\t\tif option.BlockCache%_const.BlockSize != 0 {\n\t\t\tblockNum++\n\t\t}\n\t\ttinyWAL.localCache, err = lru.New[uint32, []byte](int(blockNum))\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\tdir, err := os.ReadDir(option.DirPath)\n\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tvar segmentFileIds []int\n\tfor _, file := range dir {\n\t\tif file.IsDir() {\n\t\t\tcontinue\n\t\t}\n\t\tsegmentFileId := 0\n\t\t_, err := fmt.Scanf(file.Name(), \"%d\"+option.segmentFileExt, &segmentFileId)\n\t\tif err != nil {\n\t\t\tcontinue\n\t\t}\n\t\tsegmentFileIds = append(segmentFileIds, segmentFileId)\n\t}\n\n\tif len(segmentFileIds) == 0 {\n\t\tsegment, err := openSegmentFile(option.DirPath, option.segmentFileExt, _const.FirstSegmentFileId, tinyWAL.localCache)\n\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\ttinyWAL.activeSegment = segment\n\t} else {\n\t\tsort.Ints(segmentFileIds)\n\t\tfor i, fileId := range segmentFileIds {\n\t\t\tsegment, err := openSegmentFile(option.DirPath, option.segmentFileExt, uint32(fileId), tinyWAL.localCache)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tif i == len(segmentFileIds)-1 {\n\t\t\t\ttinyWAL.activeSegment = segment\n\t\t\t} else {\n\t\t\t\ttinyWAL.immutableSegment[uint32(fileId)] = segment\n\t\t\t}\n\t\t}\n\t}\n\n\treturn tinyWAL, nil\n}\n\nfunc (w *TinyWAL) PendingWrites(data []byte) error {\n\tw.pendingWritesLock.Lock()\n\tdefer w.pendingWritesLock.Unlock()\n\n\tw.maxWriteSize(int64(len(data)))\n\treturn nil\n}\n\nfunc (w *TinyWAL) WriteAll() ([]*ChunkPosition, error) {\n\tif len(w.pendingWrites) == 0 {\n\t\treturn make([]*ChunkPosition, 0), nil\n\t}\n\tw.mutex.Lock()\n\tdefer func() {\n\t\tw.ClearPendingWrites()\n\t\tw.mutex.Unlock()\n\t}()\n\n\tif w.pendingWritesSize > w.option.MemTableSize {\n\t\treturn nil, _const.ErrorPendingSizeTooLarge\n\t}\n\n\tif uint64(w.activeSegment.Size())+w.pendingWritesSize > w.option.MemTableSize {\n\t\terr := w.replaceActiveSegmentFile()\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\tall, err := w.activeSegment.WriteAll(w.pendingWrites)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn all, nil\n}\n\nfunc (w *TinyWAL) Sync() error {\n\tw.mutex.Lock()\n\tdefer w.mutex.Unlock()\n\n\treturn w.activeSegment.Sync()\n}\n\nfunc (w *TinyWAL) maxWriteSize(size int64) int64 {\n\tchunks := (size + _const.BlockSize - 1) / _const.BlockSize // 计算正确的块数（向上取整）\n\ttotal := chunks * _const.ChunkHeadSize                     // 总块头大小\n\tnewHeadSize := _const.ChunkHeadSize + size                 // 基础头+数据大小\n\treturn newHeadSize + total\n}\n\nfunc (w *TinyWAL) NewReader() *_const.Reader {\n\tw.mutex.RLock()\n\tdefer w.mutex.RUnlock()\n\tvar readers []*SegmentReader\n\tfor _, segment := range w.immutableSegment {\n\t\treaders = append(readers, segment.NewSegmentReader())\n\t}\n\treaders = append(readers, w.activeSegment.NewSegmentReader())\n\n\tsort.Slice(readers, func(i, j int) bool { return readers[i].seg.segmentFileId < readers[j].seg.segmentFileId })\n\n\treturn &_const.Reader{\n\t\tAllSegmentReader: readers,\n\t\tProgress:         0,\n\t}\n\n}\n\nfunc (w *TinyWAL) Write(data []byte) (*ChunkPosition, error) {\n\tw.mutex.Lock()\n\tdefer w.mutex.Unlock()\n\tif w.maxWriteSize(int64(len(data))) > int64(w.option.MemTableSize) {\n\t\treturn nil, _const.ErrorDataToLarge\n\t}\n\n\tif w.isFull(int64(len(data))) {\n\t\tif err := w.replaceActiveSegmentFile(); err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\tposition, err := w.activeSegment.Write(data)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tw.byteWrite += uint64(position.ChunkSize)\n\n\tisSync := w.option.Sync\n\tif !isSync && w.byteWrite > w.option.BytesPerSync {\n\t\tisSync = true\n\t}\n\tif isSync {\n\t\terr := w.activeSegment.Sync()\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tw.byteWrite = 0\n\t}\n\treturn position, err\n}\n\nfunc (w *TinyWAL) isFull(delta int64) bool {\n\treturn w.activeSegment.Size()+w.maxWriteSize(delta) > int64(w.option.MemTableSize)\n}\n\nfunc (w *TinyWAL) replaceActiveSegmentFile() error {\n\terr := w.activeSegment.Sync()\n\tif err != nil {\n\t\treturn err\n\t}\n\tw.byteWrite = 0\n\tfile, err := openSegmentFile(w.option.DirPath, w.option.segmentFileExt, w.activeSegment.segmentFileId+1, w.localCache)\n\tif err != nil {\n\t\treturn err\n\t}\n\tw.immutableSegment[w.activeSegment.segmentFileId] = w.activeSegment\n\tw.activeSegment = file\n\treturn nil\n}\n\nfunc (w *TinyWAL) ClearPendingWrites() {\n\n}\n"
  }
]