github.com/gnolang/gno@v0.0.0-20240520182011-228e9d0192ce/tm2/pkg/db/internal/mem_batch.go (about)

     1  package internal
     2  
     3  import "sync"
     4  
     5  type AtomicSetDeleter interface {
     6  	Mutex() *sync.Mutex
     7  	SetNoLock(key, value []byte)
     8  	SetNoLockSync(key, value []byte)
     9  	DeleteNoLock(key []byte)
    10  	DeleteNoLockSync(key []byte)
    11  }
    12  
    13  type MemBatch struct {
    14  	DB  AtomicSetDeleter
    15  	Ops []Operation
    16  }
    17  
    18  type OpType int
    19  
    20  const (
    21  	OpTypeSet    OpType = 1
    22  	OpTypeDelete OpType = 2
    23  )
    24  
    25  type Operation struct {
    26  	OpType
    27  	Key   []byte
    28  	Value []byte
    29  }
    30  
    31  func (mBatch *MemBatch) Set(key, value []byte) {
    32  	mBatch.Ops = append(mBatch.Ops, Operation{OpTypeSet, key, value})
    33  }
    34  
    35  func (mBatch *MemBatch) Delete(key []byte) {
    36  	mBatch.Ops = append(mBatch.Ops, Operation{OpTypeDelete, key, nil})
    37  }
    38  
    39  func (mBatch *MemBatch) Write() {
    40  	mBatch.write(false)
    41  }
    42  
    43  func (mBatch *MemBatch) WriteSync() {
    44  	mBatch.write(true)
    45  }
    46  
    47  func (mBatch *MemBatch) Close() {
    48  	mBatch.Ops = nil
    49  }
    50  
    51  func (mBatch *MemBatch) write(doSync bool) {
    52  	if mtx := mBatch.DB.Mutex(); mtx != nil {
    53  		mtx.Lock()
    54  		defer mtx.Unlock()
    55  	}
    56  
    57  	for i, op := range mBatch.Ops {
    58  		if doSync && i == (len(mBatch.Ops)-1) {
    59  			switch op.OpType {
    60  			case OpTypeSet:
    61  				mBatch.DB.SetNoLockSync(op.Key, op.Value)
    62  			case OpTypeDelete:
    63  				mBatch.DB.DeleteNoLockSync(op.Key)
    64  			}
    65  			break // we're done.
    66  		}
    67  		switch op.OpType {
    68  		case OpTypeSet:
    69  			mBatch.DB.SetNoLock(op.Key, op.Value)
    70  		case OpTypeDelete:
    71  			mBatch.DB.DeleteNoLock(op.Key)
    72  		}
    73  	}
    74  }