discard.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 "encoding/binary"
10 "os"
11 "path/filepath"
12 "sort"
13 "sync"
14
15 "github.com/dgraph-io/badger/v4/y"
16 "github.com/dgraph-io/ristretto/v2/z"
17 )
18
19 // discardStats keeps track of the amount of data that could be discarded for
20 // a given logfile.
21 type discardStats struct {
22 sync.Mutex
23
24 *z.MmapFile
25 opt Options
26 nextEmptySlot int
27 }
28
29 const discardFname string = "DISCARD"
30
31 func InitDiscardStats(opt Options) (*discardStats, error) {
32 fname := filepath.Join(opt.ValueDir, discardFname)
33
34 // 1MB file can store 65.536 discard entries. Each entry is 16 bytes.
35 mf, err := z.OpenMmapFile(fname, os.O_CREATE|os.O_RDWR, 1<<20)
36 lf := &discardStats{
37 MmapFile: mf,
38 opt: opt,
39 }
40 if err == z.NewFile {
41 // We don't need to zero out the entire 1MB.
42 lf.zeroOut()
43
44 } else if err != nil {
45 return nil, y.Wrapf(err, "while opening file: %s\n", discardFname)
46 }
47
48 for slot := 0; slot < lf.maxSlot(); slot++ {
49 if lf.get(16*slot) == 0 {
50 lf.nextEmptySlot = slot
51 break
52 }
53 }
54 sort.Sort(lf)
55 opt.Infof("Discard stats nextEmptySlot: %d\n", lf.nextEmptySlot)
56 return lf, nil
57 }
58
59 func (lf *discardStats) Len() int {
60 return lf.nextEmptySlot
61 }
62 func (lf *discardStats) Less(i, j int) bool {
63 return lf.get(16*i) < lf.get(16*j)
64 }
65 func (lf *discardStats) Swap(i, j int) {
66 left := lf.Data[16*i : 16*i+16]
67 right := lf.Data[16*j : 16*j+16]
68 var tmp [16]byte
69 copy(tmp[:], left)
70 copy(left, right)
71 copy(right, tmp[:])
72 }
73
74 // offset is not slot.
75 func (lf *discardStats) get(offset int) uint64 {
76 return binary.BigEndian.Uint64(lf.Data[offset : offset+8])
77 }
78 func (lf *discardStats) set(offset int, val uint64) {
79 binary.BigEndian.PutUint64(lf.Data[offset:offset+8], val)
80 }
81
82 // zeroOut would zero out the next slot.
83 func (lf *discardStats) zeroOut() {
84 lf.set(lf.nextEmptySlot*16, 0)
85 lf.set(lf.nextEmptySlot*16+8, 0)
86 }
87
88 func (lf *discardStats) maxSlot() int {
89 return len(lf.Data) / 16
90 }
91
92 // Update would update the discard stats for the given file id. If discard is
93 // 0, it would return the current value of discard for the file. If discard is
94 // < 0, it would set the current value of discard to zero for the file.
95 func (lf *discardStats) Update(fidu uint32, discard int64) int64 {
96 fid := uint64(fidu)
97 lf.Lock()
98 defer lf.Unlock()
99
100 idx := sort.Search(lf.nextEmptySlot, func(slot int) bool {
101 return lf.get(slot*16) >= fid
102 })
103 if idx < lf.nextEmptySlot && lf.get(idx*16) == fid {
104 off := idx*16 + 8
105 curDisc := lf.get(off)
106 if discard == 0 {
107 return int64(curDisc)
108 }
109 if discard < 0 {
110 lf.set(off, 0)
111 return 0
112 }
113 lf.set(off, curDisc+uint64(discard))
114 return int64(curDisc + uint64(discard))
115 }
116 if discard <= 0 {
117 // No need to add a new entry.
118 return 0
119 }
120
121 // Could not find the fid. Add the entry.
122 idx = lf.nextEmptySlot
123 lf.set(idx*16, fid)
124 lf.set(idx*16+8, uint64(discard))
125
126 // Move to next slot.
127 lf.nextEmptySlot++
128 for lf.nextEmptySlot >= lf.maxSlot() {
129 y.Check(lf.Truncate(2 * int64(len(lf.Data))))
130 }
131 lf.zeroOut()
132
133 sort.Sort(lf)
134 return discard
135 }
136
137 func (lf *discardStats) Iterate(f func(fid, stats uint64)) {
138 for slot := 0; slot < lf.nextEmptySlot; slot++ {
139 idx := 16 * slot
140 f(lf.get(idx), lf.get(idx+8))
141 }
142 }
143
144 // MaxDiscard returns the file id with maximum discard bytes.
145 func (lf *discardStats) MaxDiscard() (uint32, int64) {
146 lf.Lock()
147 defer lf.Unlock()
148
149 var maxFid, maxVal uint64
150 lf.Iterate(func(fid, val uint64) {
151 if maxVal < val {
152 maxVal = val
153 maxFid = fid
154 }
155 })
156 return uint32(maxFid), int64(maxVal)
157 }
158