steamcache2/vfs/cache/cache.go

package cache

import (
	"fmt"
	"s1d3sw1ped/SteamCache2/vfs"
	"s1d3sw1ped/SteamCache2/vfs/cachestate"
	"s1d3sw1ped/SteamCache2/vfs/vfserror"
	"sync"
)

// Ensure CacheFS implements VFS.
var _ vfs.VFS = (*CacheFS)(nil)

// CacheFS is a virtual file system that caches files in memory and on disk.
type CacheFS struct {
	fast vfs.VFS
	slow vfs.VFS

	cacheHandler CacheHandler

	keyLocks sync.Map // map[string]*sync.RWMutex for per-key locks
}

type CacheHandler func(*vfs.FileInfo, cachestate.CacheState) bool

// New creates a new CacheFS. fast is used for caching, and slow is used for storage. fast should obviously be faster than slow.
func New(cacheHandler CacheHandler) *CacheFS {
	return &CacheFS{
		cacheHandler: cacheHandler,
		keyLocks:     sync.Map{},
	}
}

func (c *CacheFS) SetSlow(vfs vfs.VFS) {
	if vfs == nil {
		panic("vfs is nil") // panic if the vfs is nil
	}

	c.slow = vfs
}

func (c *CacheFS) SetFast(vfs vfs.VFS) {
	c.fast = vfs
}

// getKeyLock returns a RWMutex for the given key, creating it if necessary.
func (c *CacheFS) getKeyLock(key string) *sync.RWMutex {
	mu, _ := c.keyLocks.LoadOrStore(key, &sync.RWMutex{})
	return mu.(*sync.RWMutex)
}

// cacheState returns the state of the file at key.
func (c *CacheFS) cacheState(key string) cachestate.CacheState {
	if c.fast != nil {
		if _, err := c.fast.Stat(key); err == nil {
			return cachestate.CacheStateHit
		}
	}

	if _, err := c.slow.Stat(key); err == nil {
		return cachestate.CacheStateMiss
	}

	return cachestate.CacheStateNotFound
}

func (c *CacheFS) Name() string {
	return fmt.Sprintf("CacheFS(%s, %s)", c.fast.Name(), c.slow.Name())
}

// Size returns the total size of the cache.
func (c *CacheFS) Size() int64 {
	return c.slow.Size()
}

// Set sets the file at key to src. If the file is already in the cache, it is replaced.
func (c *CacheFS) Set(key string, src []byte) error {
	mu := c.getKeyLock(key)
	mu.Lock()
	defer mu.Unlock()

	state := c.cacheState(key)

	switch state {
	case cachestate.CacheStateHit:
		if c.fast != nil {
			c.fast.Delete(key)
		}
		return c.slow.Set(key, src)
	case cachestate.CacheStateMiss, cachestate.CacheStateNotFound:
		return c.slow.Set(key, src)
	}

	panic(vfserror.ErrUnreachable)
}

// Delete deletes the file at key from the cache.
func (c *CacheFS) Delete(key string) error {
	mu := c.getKeyLock(key)
	mu.Lock()
	defer mu.Unlock()

	if c.fast != nil {
		c.fast.Delete(key)
	}
	return c.slow.Delete(key)
}

// Get returns the file at key. If the file is not in the cache, it is fetched from the storage.
func (c *CacheFS) Get(key string) ([]byte, error) {
	src, _, err := c.GetS(key)
	return src, err
}

// GetS returns the file at key. If the file is not in the cache, it is fetched from the storage. It also returns the cache state.
func (c *CacheFS) GetS(key string) ([]byte, cachestate.CacheState, error) {
	mu := c.getKeyLock(key)
	mu.RLock()
	defer mu.RUnlock()

	state := c.cacheState(key)

	switch state {
	case cachestate.CacheStateHit:
		// if c.fast == nil then cacheState cannot be CacheStateHit so we can safely ignore the check
		src, err := c.fast.Get(key)
		return src, state, err
	case cachestate.CacheStateMiss:
		src, err := c.slow.Get(key)
		if err != nil {
			return nil, state, err
		}

		sstat, _ := c.slow.Stat(key)
		if sstat != nil && c.fast != nil { // file found in slow storage and fast storage is available
			// We are accessing the file from the slow storage, and the file has been accessed less then a minute ago so it popular, so we should update the fast storage with the latest file.
			if c.cacheHandler != nil && c.cacheHandler(sstat, state) {
				if err := c.fast.Set(key, src); err != nil {
					return nil, state, err
				}
			}
		}

		return src, state, nil
	case cachestate.CacheStateNotFound:
		return nil, state, vfserror.ErrNotFound
	}

	panic(vfserror.ErrUnreachable)
}

// Stat returns information about the file at key.
// Warning: This will return information about the file in the fastest storage its in.
func (c *CacheFS) Stat(key string) (*vfs.FileInfo, error) {
	mu := c.getKeyLock(key)
	mu.RLock()
	defer mu.RUnlock()

	state := c.cacheState(key)

	switch state {
	case cachestate.CacheStateHit:
		// if c.fast == nil then cacheState cannot be CacheStateHit so we can safely ignore the check
		return c.fast.Stat(key)
	case cachestate.CacheStateMiss:
		return c.slow.Stat(key)
	case cachestate.CacheStateNotFound:
		return nil, vfserror.ErrNotFound
	}

	panic(vfserror.ErrUnreachable)
}

// StatAll returns information about all files in the cache.
// Warning: This only returns information about the files in the slow storage.
func (c *CacheFS) StatAll() []*vfs.FileInfo {
	return c.slow.StatAll()
}