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 }