package writecachebadger import ( "encoding/binary" "errors" "sync" "git.frostfs.info/TrueCloudLab/frostfs-node/internal/logs" "github.com/dgraph-io/badger/v4" "go.uber.org/atomic" "go.uber.org/zap" ) var ( errCountCancelled = errors.New("counting was cancelled") ) func (c *cache) estimateCacheSize() uint64 { onDiskSize, _ := c.db.EstimateSize(nil) c.metrics.SetEstimateSize(onDiskSize, 0) return onDiskSize } func (c *cache) incSizeDB(sz uint64) uint64 { return sz + c.maxObjectSize } type countAsync struct { value atomic.Uint64 // mutable checks whether the value can be incremented/decremented. mutable bool // mutMtx is mutex for mutable flag. mutMtx sync.RWMutex // init checks whether the value is finally initialized. init bool // initMtx is mutex for init flag. initMtx sync.RWMutex // modeCh is channel to indicate SetMode has been invoked and // counting must be cancelled. modeCh chan struct{} // modeDoneOnce guarantees that modeCh has been closed only once. // This may happen if SetMode has been invoked many times. modeDoneOnce sync.Once // closeCh is channel to indicate Close has been invoked and // counting must be cancelled. closeCh chan struct{} // closeDoneOnce guarantees that modeCh has been closed only once. // This may happen if Close has been invoked many times. closeDoneOnce sync.Once } func (c *countAsync) cancelByClose() { c.closeDoneOnce.Do(func() { close(c.closeCh) }) } func (c *countAsync) cancelBySetMode() { c.modeDoneOnce.Do(func() { close(c.modeCh) }) } func (c *countAsync) addAndSetAsInit(value uint64) { c.value.Add(value) c.initMtx.Lock() defer c.initMtx.Unlock() c.init = true } func (c *countAsync) isInitialized() bool { c.initMtx.RLock() defer c.initMtx.RUnlock() return c.init } func (c *countAsync) isMutable() bool { c.mutMtx.RLock() defer c.mutMtx.RUnlock() return c.mutable } func (c *countAsync) setAsMutable() { c.mutMtx.Lock() defer c.mutMtx.Unlock() c.mutable = true } func (c *countAsync) reset() { c.value.Store(0) c.closeCh = make(chan struct{}) c.closeDoneOnce = sync.Once{} c.modeCh = make(chan struct{}) c.modeDoneOnce = sync.Once{} c.initMtx.Lock() defer c.initMtx.Unlock() c.init = false c.mutMtx.Lock() defer c.mutMtx.Unlock() c.mutable = false } type counters struct { cDB countAsync } func (x *counters) IncDB() { if x.cDB.isMutable() { x.cDB.value.Inc() } } func (x *counters) DecDB() { if x.cDB.isMutable() { x.cDB.value.Dec() } } func (x *counters) DB() uint64 { if x.cDB.isInitialized() { return x.cDB.value.Load() } return 0 } func (x *counters) isReadyToFlush() bool { return x.cDB.isInitialized() } func (c *cache) initCounters() error { c.objCounters.cDB.reset() c.initCounterWG.Add(1) go func() { defer func() { c.initCounterWG.Done() }() var inDB uint64 if !c.options.recountKeys { k := keyCountPrefix cDB, err := c.getRawKey(k[:]) if err == nil { c.objCounters.cDB.setAsMutable() inDB = binary.LittleEndian.Uint64(cDB) c.modeMtx.RLock() defer c.modeMtx.RUnlock() if !c.mode.ReadOnly() { if err = c.deleteRawKey(k[:]); err != nil { c.log.Error(logs.WritecacheBadgerCouldNotDeleteKeysCount, zap.Error(err)) } } return } } // Recount keys if "key_count" has not been found if err := c.db.View(func(tx *badger.Txn) error { c.objCounters.cDB.setAsMutable() opts := badger.DefaultIteratorOptions opts.PrefetchValues = false it := tx.NewIterator(opts) defer it.Close() for it.Rewind(); it.Valid(); it.Next() { select { case <-c.objCounters.cDB.closeCh: return errCountCancelled case <-c.objCounters.cDB.modeCh: return errCountCancelled default: inDB++ } } return nil }); err != nil { c.log.Error(logs.WritecacheBadgerCouldNotCountKeys, zap.Error(err)) return } c.metrics.SetActualCounters(inDB, 0) c.objCounters.cDB.addAndSetAsInit(inDB) }() return nil }