github.com/nutsdb/nutsdb@v1.0.4/batch.go (about)

     1  package nutsdb
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"sync"
     7  	"sync/atomic"
     8  )
     9  
    10  // ErrCommitAfterFinish indicates that write batch commit was called after
    11  var ErrCommitAfterFinish = errors.New("Batch commit not permitted after finish")
    12  
    13  const (
    14  	DefaultThrottleSize = 16
    15  )
    16  
    17  // WriteBatch holds the necessary info to perform batched writes.
    18  type WriteBatch struct {
    19  	sync.Mutex
    20  	tx       *Tx
    21  	db       *DB
    22  	throttle *Throttle
    23  	err      atomic.Value
    24  	finished bool
    25  }
    26  
    27  func (db *DB) NewWriteBatch() (*WriteBatch, error) {
    28  	wb := &WriteBatch{
    29  		db:       db,
    30  		throttle: NewThrottle(DefaultThrottleSize),
    31  	}
    32  
    33  	var err error
    34  	wb.tx, err = newTx(db, true)
    35  	return wb, err
    36  }
    37  
    38  // SetMaxPendingTxns sets a limit on maximum number of pending transactions while writing batches.
    39  // This function should be called before using WriteBatch. Default value of MaxPendingTxns is
    40  // 16 to minimise memory usage.
    41  func (wb *WriteBatch) SetMaxPendingTxns(max int) {
    42  	wb.throttle = NewThrottle(max)
    43  }
    44  
    45  func (wb *WriteBatch) Cancel() error {
    46  	wb.Lock()
    47  	wb.finished = true
    48  	wb.tx.setStatusClosed()
    49  	wb.Unlock()
    50  
    51  	if err := wb.throttle.Finish(); err != nil {
    52  		return fmt.Errorf("WatchBatch.Cancel error while finishing: %v", err)
    53  	}
    54  
    55  	return nil
    56  }
    57  
    58  func (wb *WriteBatch) Put(bucket string, key, value []byte, ttl uint32) error {
    59  	wb.Lock()
    60  	defer wb.Unlock()
    61  
    62  	wb.tx.lock()
    63  	err := wb.tx.Put(bucket, key, value, ttl)
    64  	if err != nil {
    65  		wb.tx.unlock()
    66  		return err
    67  	}
    68  
    69  	if nil == wb.tx.checkSize() {
    70  		wb.tx.unlock()
    71  		return err
    72  	}
    73  
    74  	wb.tx.unlock()
    75  	if cerr := wb.commit(); cerr != nil {
    76  		return cerr
    77  	}
    78  
    79  	return err
    80  }
    81  
    82  // func (tx *Tx) Delete(bucket string, key []byte) error
    83  func (wb *WriteBatch) Delete(bucket string, key []byte) error {
    84  	wb.Lock()
    85  	defer wb.Unlock()
    86  
    87  	wb.tx.lock()
    88  	err := wb.tx.Delete(bucket, key)
    89  	if err != nil {
    90  		wb.tx.unlock()
    91  		return err
    92  	}
    93  
    94  	if nil == wb.tx.checkSize() {
    95  		wb.tx.unlock()
    96  		return err
    97  	}
    98  
    99  	wb.tx.unlock()
   100  	if cerr := wb.commit(); cerr != nil {
   101  		return cerr
   102  	}
   103  
   104  	return err
   105  }
   106  
   107  func (wb *WriteBatch) commit() error {
   108  	if err := wb.Error(); err != nil {
   109  		return err
   110  	}
   111  
   112  	if wb.finished {
   113  		return ErrCommitAfterFinish
   114  	}
   115  
   116  	if err := wb.throttle.Do(); err != nil {
   117  		wb.err.Store(err)
   118  		return err
   119  	}
   120  
   121  	wb.tx.CommitWith(wb.callback)
   122  
   123  	// new a new tx
   124  	var err error
   125  	wb.tx, err = newTx(wb.db, true)
   126  	if err != nil {
   127  		return err
   128  	}
   129  
   130  	return wb.Error()
   131  }
   132  
   133  func (wb *WriteBatch) callback(err error) {
   134  	// sync.WaitGroup is thread-safe, so it doesn't need to be run inside wb.Lock.
   135  	defer wb.throttle.Done(err)
   136  	if err == nil {
   137  		return
   138  	}
   139  	if err := wb.Error(); err != nil {
   140  		return
   141  	}
   142  
   143  	wb.err.Store(err)
   144  }
   145  
   146  func (wb *WriteBatch) Flush() error {
   147  	wb.Lock()
   148  	err := wb.commit()
   149  	if err != nil {
   150  		wb.Unlock()
   151  		return err
   152  	}
   153  	wb.finished = true
   154  	wb.tx.setStatusClosed()
   155  	wb.Unlock()
   156  	if err := wb.throttle.Finish(); err != nil {
   157  		if wb.Error() != nil {
   158  			return fmt.Errorf("wb.err: %s err: %s", wb.Error(), err)
   159  		}
   160  		return err
   161  	}
   162  
   163  	return wb.Error()
   164  }
   165  
   166  func (wb *WriteBatch) Reset() error {
   167  	wb.Lock()
   168  	defer wb.Unlock()
   169  	var err error
   170  	wb.finished = false
   171  	wb.tx, err = newTx(wb.db, true)
   172  	if err != nil {
   173  		return err
   174  	}
   175  	wb.throttle = NewThrottle(DefaultThrottleSize)
   176  	return err
   177  }
   178  
   179  // Error returns any errors encountered so far. No commits would be run once an error is detected.
   180  func (wb *WriteBatch) Error() error {
   181  	// If the interface conversion fails, the err will be nil.
   182  	err, _ := wb.err.Load().(error)
   183  	return err
   184  }