github.com/blong14/gache@v0.0.0-20240124023949-89416fd8bbfa/internal/db/memtable/table.go (about)

     1  package memtable
     2  
     3  import (
     4  	"errors"
     5  	"sync/atomic"
     6  	"unsafe"
     7  
     8  	gstable "github.com/blong14/gache/internal/db/sstable"
     9  )
    10  
    11  var ErrAllowedBytesExceeded = errors.New("memtable allowed bytes exceeded")
    12  
    13  type MemTable struct {
    14  	readBuffer *SkipList
    15  	bytes      uint64
    16  }
    17  
    18  func New() *MemTable {
    19  	return &MemTable{
    20  		readBuffer: NewSkipList(),
    21  	}
    22  }
    23  
    24  func (m *MemTable) buffer() *SkipList {
    25  	reader := (*SkipList)(atomic.LoadPointer(
    26  		(*unsafe.Pointer)(unsafe.Pointer(&m.readBuffer))))
    27  	return reader
    28  }
    29  
    30  func (m *MemTable) Get(k []byte) ([]byte, bool) {
    31  	return m.buffer().Get(k)
    32  }
    33  
    34  func (m *MemTable) Count() uint64 {
    35  	return m.buffer().Count()
    36  }
    37  
    38  func (m *MemTable) Set(k, v []byte) error {
    39  	err := m.buffer().Set(k, v)
    40  	if err != nil {
    41  		return err
    42  	}
    43  	byts := atomic.AddUint64(&m.bytes, uint64(len(k)+len(v)))
    44  	if byts >= 4096*4096 {
    45  		return ErrAllowedBytesExceeded
    46  	}
    47  	return nil
    48  }
    49  
    50  func (m *MemTable) Scan(k, v []byte, f func(k, v []byte) bool) {
    51  	m.buffer().Scan(k, v, f)
    52  }
    53  
    54  func (m *MemTable) Range(f func(k, v []byte) bool) {
    55  	m.buffer().Range(f)
    56  }
    57  
    58  func (m *MemTable) Flush(sstable *gstable.SSTable) error {
    59  	reader := m.buffer()
    60  	nReader := NewSkipList()
    61  	for {
    62  		if atomic.CompareAndSwapPointer(
    63  			(*unsafe.Pointer)(unsafe.Pointer(&m.readBuffer)),
    64  			unsafe.Pointer(reader),
    65  			unsafe.Pointer(nReader),
    66  		) {
    67  			atomic.StoreUint64(&m.bytes, 0)
    68  			reader.Range(func(k, v []byte) bool {
    69  				err := sstable.Set(k, v)
    70  				return err == nil
    71  			})
    72  			return nil
    73  		}
    74  		reader = m.buffer()
    75  	}
    76  }