frostfs-node/pkg/local_object_storage/blobstor/blobovniczatree/cache.go

package blobovniczatree

import (
	"context"
	"sync"
	"time"

	utilSync "git.frostfs.info/TrueCloudLab/frostfs-node/pkg/util/sync"
	cache "github.com/go-pkgz/expirable-cache/v3"
)

// dbCache caches sharedDB instances that are NOT open for Put.
//
// Uses dbManager for opening/closing sharedDB instances.
// Stores a reference to an cached sharedDB, so dbManager does not close it.
type dbCache struct {
	cacheGuard *sync.Mutex
	cache      cache.Cache[string, *sharedDB]
	pathLock   *utilSync.KeyLocker[string] // the order of locks is important: pathLock first, cacheGuard second
	closed     bool
	nonCached  map[string]struct{}
	wg         sync.WaitGroup
	cancel     context.CancelFunc

	dbManager *dbManager
}

func newDBCache(parentCtx context.Context, size int, ttl time.Duration, dbManager *dbManager) *dbCache {
	ch := cache.NewCache[string, *sharedDB]().
		WithTTL(ttl).WithLRU().WithMaxKeys(size).
		WithOnEvicted(func(_ string, db *sharedDB) {
			db.Close()
		})
	ctx, cancel := context.WithCancel(parentCtx)
	res := &dbCache{
		cacheGuard: &sync.Mutex{},
		wg:         sync.WaitGroup{},
		cancel:     cancel,
		cache:      ch,
		dbManager:  dbManager,
		pathLock:   utilSync.NewKeyLocker[string](),
		nonCached:  make(map[string]struct{}),
	}
	if ttl > 0 {
		res.wg.Add(1)
		go func() {
			ticker := time.NewTicker(ttl / 100)
			defer ticker.Stop()
			for {
				select {
				case <-ctx.Done():
					res.wg.Done()
					return
				case <-ticker.C:
					res.cacheGuard.Lock()
					res.cache.DeleteExpired()
					res.cacheGuard.Unlock()
				}
			}
		}()
	}
	return res
}

func (c *dbCache) Open() {
	c.cacheGuard.Lock()
	defer c.cacheGuard.Unlock()

	c.closed = false
}

func (c *dbCache) Close() {
	c.cacheGuard.Lock()
	defer c.cacheGuard.Unlock()
	c.cancel()
	c.wg.Wait()
	c.cache.Purge()
	c.closed = true
}

func (c *dbCache) GetOrCreate(path string) *sharedDB {
	value := c.getExisted(path)
	if value != nil {
		return value
	}
	return c.create(path)
}

func (c *dbCache) EvictAndMarkNonCached(path string) {
	c.pathLock.Lock(path)
	defer c.pathLock.Unlock(path)

	c.cacheGuard.Lock()
	defer c.cacheGuard.Unlock()

	c.cache.Remove(path)
	c.nonCached[path] = struct{}{}
}

func (c *dbCache) RemoveFromNonCached(path string) {
	c.pathLock.Lock(path)
	defer c.pathLock.Unlock(path)

	c.cacheGuard.Lock()
	defer c.cacheGuard.Unlock()

	delete(c.nonCached, path)
}

func (c *dbCache) getExisted(path string) *sharedDB {
	c.cacheGuard.Lock()
	defer c.cacheGuard.Unlock()

	if value, ok := c.cache.Get(path); ok {
		return value
	} else if value != nil {
		c.cache.Invalidate(path)
	}
	return nil
}

func (c *dbCache) create(path string) *sharedDB {
	c.pathLock.Lock(path)
	defer c.pathLock.Unlock(path)

	value := c.getExisted(path)
	if value != nil {
		return value
	}

	value = c.dbManager.GetByPath(path)

	_, err := value.Open() // open db to hold reference, closed by evictedDB.Close() or if cache closed
	if err != nil {
		return value
	}
	if added := c.put(path, value); !added {
		value.Close()
	}
	return value
}

func (c *dbCache) put(path string, db *sharedDB) bool {
	c.cacheGuard.Lock()
	defer c.cacheGuard.Unlock()

	_, isNonCached := c.nonCached[path]

	if isNonCached || c.closed {
		return false
	}
	c.cache.Add(path, db)
	return true
}