steamcache2/vfs/cache/cache.go

// vfs/cache/cache.go
package cache

import (
	"fmt"
	"io"
	"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()
}

// 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)
}

// Open returns the file at key. If the file is not in the cache, it is fetched from the storage.
func (c *CacheFS) Open(key string) (io.ReadCloser, 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.Open(key)
	case cachestate.CacheStateMiss:
		slowReader, err := c.slow.Open(key)
		if err != nil {
			return nil, 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) {
				fastWriter, err := c.fast.Create(key, sstat.Size())
				if err == nil {
					return &teeReadCloser{
						Reader:  io.TeeReader(slowReader, fastWriter),
						closers: []io.Closer{slowReader, fastWriter},
					}, nil
				}
			}
		}

		return slowReader, nil
	case cachestate.CacheStateNotFound:
		return nil, vfserror.ErrNotFound
	}

	panic(vfserror.ErrUnreachable)
}

// Create creates a new file at key. If the file is already in the cache, it is replaced.
func (c *CacheFS) Create(key string, size int64) (io.WriteCloser, 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.Create(key, size)
	case cachestate.CacheStateMiss, cachestate.CacheStateNotFound:
		return c.slow.Create(key, size)
	}

	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()
}

type teeReadCloser struct {
	io.Reader
	closers []io.Closer
}

func (t *teeReadCloser) Close() error {
	var err error
	for _, c := range t.closers {
		if e := c.Close(); e != nil {
			err = e
		}
	}
	return err
}