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

     1  package memtable
     2  
     3  import (
     4  	"errors"
     5  	"github.com/scottcagno/storage/pkg/_junk/_lsmtree/container/rbtree"
     6  	"github.com/scottcagno/storage/pkg/_junk/_lsmtree/sstable"
     7  	"github.com/scottcagno/storage/pkg/_junk/_x/file"
     8  	"os"
     9  	"runtime"
    10  )
    11  
    12  var ErrNotFound = errors.New("error: value not found")
    13  
    14  type Memtable struct {
    15  	base string // base is the base path of the db
    16  	rbt  *rbtree.RBTree
    17  	//wal  *wal.WAL
    18  	wal *file.SegmentManager
    19  }
    20  
    21  func Open(base string) (*Memtable, error) {
    22  	wall, err := file.Open(base)
    23  	if err != nil {
    24  		return nil, err
    25  	}
    26  	mem := &Memtable{
    27  		base: base,
    28  		rbt:  rbtree.NewRBTree(),
    29  		wal:  wall,
    30  	}
    31  	return mem, nil
    32  }
    33  
    34  func (m *Memtable) Size() int64 {
    35  	return m.rbt.Size()
    36  }
    37  
    38  func (m *Memtable) Put(key string, value []byte) (int64, error) {
    39  	// write entry to the wal
    40  	_, err := m.wal.Write(key, value)
    41  	if err != nil {
    42  		return -1, err
    43  	}
    44  	// write entry to the memtable
    45  	_, _ = m.rbt.Put(key, value)
    46  	// success, return size
    47  	return m.rbt.Size(), nil
    48  }
    49  
    50  func (m *Memtable) Get(key string) ([]byte, error) {
    51  	// check memtable for value
    52  	value, ok := m.rbt.Get(key)
    53  	if !ok {
    54  		return nil, ErrNotFound
    55  	}
    56  	// return value
    57  	return value, nil
    58  }
    59  
    60  func (m *Memtable) Del(key string) (int64, error) {
    61  	// write del entry to the wal
    62  	_, err := m.wal.Write(key, nil)
    63  	if err != nil {
    64  		return -1, err
    65  	}
    66  	// write del entry to the memtable
    67  	_, _ = m.rbt.Put(key, nil)
    68  	// return size
    69  	return m.rbt.Size(), nil
    70  }
    71  
    72  func (m *Memtable) FlushToSSTableBatch(batch *sstable.Batch) {
    73  	m.rbt.Scan(func(key string, value []byte) bool {
    74  		batch.Write(key, value)
    75  		return true
    76  	})
    77  }
    78  
    79  func (m *Memtable) Scan(iter func(key string, value []byte) bool) {
    80  	m.rbt.Scan(iter)
    81  }
    82  
    83  func (m *Memtable) Reset() error {
    84  	walPath := m.wal.Path()
    85  	err := m.wal.Close()
    86  	if err != nil {
    87  		return err
    88  	}
    89  	err = os.RemoveAll(walPath)
    90  	if err != nil {
    91  		return err
    92  	}
    93  	m.wal, err = file.Open(m.base)
    94  	if err != nil {
    95  		return err
    96  	}
    97  	m.rbt.Close()
    98  	runtime.GC()
    99  	m.rbt = rbtree.NewRBTree()
   100  	return nil
   101  }
   102  
   103  func (m *Memtable) Close() error {
   104  	err := m.wal.Close()
   105  	if err != nil {
   106  		return err
   107  	}
   108  	return nil
   109  }