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 }