frostfs-node/pkg/local_object_storage/writecache/writecachebadger/state.go

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
}