github.com/scottcagno/storage@v1.8.0/pkg/_junk/_lsmtree/db.go (about)

     1  package lsmtree
     2  
     3  import (
     4  	"github.com/scottcagno/storage/pkg/_junk/_lsmtree/memtable"
     5  	"github.com/scottcagno/storage/pkg/_junk/_lsmtree/sstable"
     6  	"log"
     7  	"path/filepath"
     8  	"sync"
     9  	"time"
    10  )
    11  
    12  const (
    13  	MaxMemtableSize      = 64 << 10 // 64 KB
    14  	defaultCommitLogPath = "wal"
    15  	defaultSSTablePath   = "data"
    16  
    17  	VERBOSE = true
    18  )
    19  
    20  type DB struct {
    21  	lock     sync.RWMutex
    22  	base     string // base is the base path of the db
    23  	mem      *memtable.Memtable
    24  	ssm      *sstable.SSManager
    25  	sstindex int64 // sst index
    26  }
    27  
    28  func Open(base string) (*DB, error) {
    29  	mem, err := memtable.Open(filepath.Join(base, defaultCommitLogPath))
    30  	if err != nil {
    31  		return nil, err
    32  	}
    33  	ssm, err := sstable.OpenSSManager(filepath.Join(base, defaultSSTablePath))
    34  	if err != nil {
    35  		return nil, err
    36  	}
    37  	db := &DB{
    38  		base:     base,
    39  		mem:      mem,
    40  		ssm:      ssm,
    41  		sstindex: ssm.GetLatestIndex(),
    42  	}
    43  	return db, nil
    44  }
    45  
    46  func (db *DB) Put(key string, value []byte) error {
    47  	// lock
    48  	db.lock.Lock()
    49  	defer db.lock.Unlock()
    50  	// pass key and value to internal upsert
    51  	return db.upsert(key, value)
    52  }
    53  
    54  func (db *DB) writeMemtableToBatch() (*sstable.Batch, error) {
    55  	// create new sstable batch to dump to
    56  	data := sstable.NewBatch()
    57  	// scan "inactive" memtable, add entries to batch
    58  	db.mem.FlushToSSTableBatch(data)
    59  	// clear memtable
    60  	err := db.mem.Reset()
    61  	if err != nil {
    62  		return nil, err
    63  	}
    64  	return data, nil
    65  }
    66  
    67  func (db *DB) upsert(key string, value []byte) error {
    68  	// insert into the memtable
    69  	size, err := db.mem.Put(key, value)
    70  	if err != nil {
    71  		return err
    72  	}
    73  	// check size
    74  	if size >= MaxMemtableSize {
    75  
    76  		log.Printf("Max Memtable size has been reached (%d KB)\nFlushing to SSTable... ", MaxMemtableSize>>10)
    77  		ts := time.Now()
    78  
    79  		// create new sstable batch
    80  		batch, err := db.writeMemtableToBatch()
    81  		if err != nil {
    82  			return err
    83  		}
    84  		// create new sstable
    85  		sst, err := sstable.CreateSSTable(filepath.Join(db.base, defaultSSTablePath), db.sstindex+1)
    86  		if err != nil {
    87  			return err
    88  		}
    89  		// write batch to sstable
    90  		err = sst.WriteBatch(batch)
    91  		if err != nil {
    92  			return err
    93  		}
    94  		// sync and close sstable
    95  		err = sst.Close()
    96  		if err != nil {
    97  			return err
    98  		}
    99  		// increment db sst index
   100  		db.sstindex++
   101  
   102  		log.Printf("[took %dms]\n", time.Since(ts).Milliseconds())
   103  	}
   104  	return nil
   105  }
   106  
   107  func (db *DB) Get(key string) ([]byte, error) {
   108  	// lock
   109  	db.lock.Lock()
   110  	defer db.lock.Unlock()
   111  	// search memtable
   112  	value, err := db.mem.Get(key)
   113  	log.Println(">>>1", err)
   114  	if err == nil {
   115  		log.Printf(">>> DEBUG-1\n")
   116  		// we found it!
   117  		log.Printf("Found value for (%s) in Memtable...\n", key)
   118  		return value, nil
   119  	}
   120  	// search sstable(s)
   121  	value, err = db.ssm.Get(key)
   122  	log.Println(">>>2", err)
   123  	if err == nil {
   124  		log.Printf(">>> DEBUG-2\n")
   125  		// we found it
   126  		log.Printf("Found value for (%s) in SSTable...\n", key)
   127  		return value, nil
   128  	}
   129  	log.Printf(">>> DEBUG-3\n")
   130  	// not fount
   131  	return nil, memtable.ErrNotFound
   132  }
   133  
   134  func (db *DB) Del(key string) error {
   135  	// lock
   136  	db.lock.Lock()
   137  	defer db.lock.Unlock()
   138  	return db.upsert(key, nil)
   139  }
   140  
   141  func (db *DB) Close() error {
   142  	err := db.mem.Close()
   143  	if err != nil {
   144  		return err
   145  	}
   146  	err = db.ssm.Close()
   147  	if err != nil {
   148  		return err
   149  	}
   150  	return nil
   151  }