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