batch.go raw
1 /*
2 * SPDX-FileCopyrightText: © Hypermode Inc. <hello@hypermode.com>
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6 package badger
7
8 import (
9 "errors"
10 "fmt"
11 "sync"
12 "sync/atomic"
13
14 "google.golang.org/protobuf/proto"
15
16 "github.com/dgraph-io/badger/v4/pb"
17 "github.com/dgraph-io/badger/v4/y"
18 "github.com/dgraph-io/ristretto/v2/z"
19 )
20
21 // WriteBatch holds the necessary info to perform batched writes.
22 type WriteBatch struct {
23 sync.Mutex
24 txn *Txn
25 db *DB
26 throttle *y.Throttle
27 err atomic.Value
28
29 isManaged bool
30 commitTs uint64
31 finished bool
32 }
33
34 // NewWriteBatch creates a new WriteBatch. This provides a way to conveniently do a lot of writes,
35 // batching them up as tightly as possible in a single transaction and using callbacks to avoid
36 // waiting for them to commit, thus achieving good performance. This API hides away the logic of
37 // creating and committing transactions. Due to the nature of SSI guaratees provided by Badger,
38 // blind writes can never encounter transaction conflicts (ErrConflict).
39 func (db *DB) NewWriteBatch() *WriteBatch {
40 if db.opt.managedTxns {
41 panic("cannot use NewWriteBatch in managed mode. Use NewWriteBatchAt instead")
42 }
43 return db.newWriteBatch(false)
44 }
45
46 func (db *DB) newWriteBatch(isManaged bool) *WriteBatch {
47 return &WriteBatch{
48 db: db,
49 isManaged: isManaged,
50 txn: db.newTransaction(true, isManaged),
51 throttle: y.NewThrottle(16),
52 }
53 }
54
55 // SetMaxPendingTxns sets a limit on maximum number of pending transactions while writing batches.
56 // This function should be called before using WriteBatch. Default value of MaxPendingTxns is
57 // 16 to minimise memory usage.
58 func (wb *WriteBatch) SetMaxPendingTxns(max int) {
59 wb.throttle = y.NewThrottle(max)
60 }
61
62 // Cancel function must be called if there's a chance that Flush might not get
63 // called. If neither Flush or Cancel is called, the transaction oracle would
64 // never get a chance to clear out the row commit timestamp map, thus causing an
65 // unbounded memory consumption. Typically, you can call Cancel as a defer
66 // statement right after NewWriteBatch is called.
67 //
68 // Note that any committed writes would still go through despite calling Cancel.
69 func (wb *WriteBatch) Cancel() {
70 wb.Lock()
71 defer wb.Unlock()
72 wb.finished = true
73 if err := wb.throttle.Finish(); err != nil {
74 wb.db.opt.Errorf("WatchBatch.Cancel error while finishing: %v", err)
75 }
76 wb.txn.Discard()
77 }
78
79 func (wb *WriteBatch) callback(err error) {
80 // sync.WaitGroup is thread-safe, so it doesn't need to be run inside wb.Lock.
81 defer wb.throttle.Done(err)
82 if err == nil {
83 return
84 }
85 if err := wb.Error(); err != nil {
86 return
87 }
88 wb.err.Store(err)
89 }
90
91 func (wb *WriteBatch) writeKV(kv *pb.KV) error {
92 e := Entry{Key: kv.Key, Value: kv.Value}
93 if len(kv.UserMeta) > 0 {
94 e.UserMeta = kv.UserMeta[0]
95 }
96 y.AssertTrue(kv.Version != 0)
97 e.version = kv.Version
98 return wb.handleEntry(&e)
99 }
100
101 func (wb *WriteBatch) Write(buf *z.Buffer) error {
102 wb.Lock()
103 defer wb.Unlock()
104
105 err := buf.SliceIterate(func(s []byte) error {
106 kv := &pb.KV{}
107 if err := proto.Unmarshal(s, kv); err != nil {
108 return err
109 }
110 return wb.writeKV(kv)
111 })
112 return err
113 }
114
115 func (wb *WriteBatch) WriteList(kvList *pb.KVList) error {
116 wb.Lock()
117 defer wb.Unlock()
118 for _, kv := range kvList.Kv {
119 if err := wb.writeKV(kv); err != nil {
120 return err
121 }
122 }
123 return nil
124 }
125
126 // SetEntryAt is the equivalent of Txn.SetEntry but it also allows setting version for the entry.
127 // SetEntryAt can be used only in managed mode.
128 func (wb *WriteBatch) SetEntryAt(e *Entry, ts uint64) error {
129 if !wb.db.opt.managedTxns {
130 return errors.New("SetEntryAt can only be used in managed mode. Use SetEntry instead")
131 }
132 e.version = ts
133 return wb.SetEntry(e)
134 }
135
136 // Should be called with lock acquired.
137 func (wb *WriteBatch) handleEntry(e *Entry) error {
138 if err := wb.txn.SetEntry(e); err != ErrTxnTooBig {
139 return err
140 }
141 // Txn has reached it's zenith. Commit now.
142 if cerr := wb.commit(); cerr != nil {
143 return cerr
144 }
145 // This time the error must not be ErrTxnTooBig, otherwise, we make the
146 // error permanent.
147 if err := wb.txn.SetEntry(e); err != nil {
148 wb.err.Store(err)
149 return err
150 }
151 return nil
152 }
153
154 // SetEntry is the equivalent of Txn.SetEntry.
155 func (wb *WriteBatch) SetEntry(e *Entry) error {
156 wb.Lock()
157 defer wb.Unlock()
158 return wb.handleEntry(e)
159 }
160
161 // Set is equivalent of Txn.Set().
162 func (wb *WriteBatch) Set(k, v []byte) error {
163 e := &Entry{Key: k, Value: v}
164 return wb.SetEntry(e)
165 }
166
167 // DeleteAt is equivalent of Txn.Delete but accepts a delete timestamp.
168 func (wb *WriteBatch) DeleteAt(k []byte, ts uint64) error {
169 e := Entry{Key: k, meta: bitDelete, version: ts}
170 return wb.SetEntry(&e)
171 }
172
173 // Delete is equivalent of Txn.Delete.
174 func (wb *WriteBatch) Delete(k []byte) error {
175 wb.Lock()
176 defer wb.Unlock()
177
178 if err := wb.txn.Delete(k); err != ErrTxnTooBig {
179 return err
180 }
181 if err := wb.commit(); err != nil {
182 return err
183 }
184 if err := wb.txn.Delete(k); err != nil {
185 wb.err.Store(err)
186 return err
187 }
188 return nil
189 }
190
191 // Caller to commit must hold a write lock.
192 func (wb *WriteBatch) commit() error {
193 if err := wb.Error(); err != nil {
194 return err
195 }
196 if wb.finished {
197 return y.ErrCommitAfterFinish
198 }
199 if err := wb.throttle.Do(); err != nil {
200 wb.err.Store(err)
201 return err
202 }
203 wb.txn.CommitWith(wb.callback)
204 wb.txn = wb.db.newTransaction(true, wb.isManaged)
205 wb.txn.commitTs = wb.commitTs
206 return wb.Error()
207 }
208
209 // Flush must be called at the end to ensure that any pending writes get committed to Badger. Flush
210 // returns any error stored by WriteBatch.
211 func (wb *WriteBatch) Flush() error {
212 wb.Lock()
213 err := wb.commit()
214 if err != nil {
215 wb.Unlock()
216 return err
217 }
218 wb.finished = true
219 wb.txn.Discard()
220 wb.Unlock()
221
222 if err := wb.throttle.Finish(); err != nil {
223 if wb.Error() != nil {
224 return fmt.Errorf("wb.err: %w err: %w", wb.Error(), err)
225 }
226 return err
227 }
228
229 return wb.Error()
230 }
231
232 // Error returns any errors encountered so far. No commits would be run once an error is detected.
233 func (wb *WriteBatch) Error() error {
234 // If the interface conversion fails, the err will be nil.
235 err, _ := wb.err.Load().(error)
236 return err
237 }
238