steamcache2/vfs/memory/memory.go

// vfs/memory/memory.go
package memory

import (
	"bytes"
	"container/list"
	"io"
	"s1d3sw1ped/SteamCache2/steamcache/logger"
	"s1d3sw1ped/SteamCache2/vfs"
	"s1d3sw1ped/SteamCache2/vfs/vfserror"
	"sync"
	"time"

	"github.com/docker/go-units"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

var (
	memoryCapacityBytes = promauto.NewGauge(
		prometheus.GaugeOpts{
			Name: "memory_cache_capacity_bytes",
			Help: "Total capacity of the memory cache in bytes",
		},
	)

	memorySizeBytes = promauto.NewGauge(
		prometheus.GaugeOpts{
			Name: "memory_cache_size_bytes",
			Help: "Total size of the memory cache in bytes",
		},
	)

	memoryReadBytes = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "memory_cache_read_bytes_total",
			Help: "Total number of bytes read from the memory cache",
		},
	)

	memoryWriteBytes = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "memory_cache_write_bytes_total",
			Help: "Total number of bytes written to the memory cache",
		},
	)
)

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

// file represents a file in memory.
type file struct {
	fileinfo *vfs.FileInfo
	data     []byte
}

// MemoryFS is a virtual file system that stores files in memory.
type MemoryFS struct {
	files    map[string]*file
	capacity int64
	size     int64
	mu       sync.RWMutex
	keyLocks sync.Map // map[string]*sync.RWMutex
	LRU      *lruList
}

// lruList for LRU eviction
type lruList struct {
	list *list.List
	elem map[string]*list.Element
}

func newLruList() *lruList {
	return &lruList{
		list: list.New(),
		elem: make(map[string]*list.Element),
	}
}

func (l *lruList) MoveToFront(key string) {
	if e, ok := l.elem[key]; ok {
		l.list.MoveToFront(e)
	}
}

func (l *lruList) Add(key string, fi *vfs.FileInfo) *list.Element {
	e := l.list.PushFront(fi)
	l.elem[key] = e
	return e
}

func (l *lruList) Remove(key string) {
	if e, ok := l.elem[key]; ok {
		l.list.Remove(e)
		delete(l.elem, key)
	}
}

func (l *lruList) Back() *vfs.FileInfo {
	if e := l.list.Back(); e != nil {
		return e.Value.(*vfs.FileInfo)
	}
	return nil
}

// New creates a new MemoryFS.
func New(capacity int64) *MemoryFS {
	if capacity <= 0 {
		panic("memory capacity must be greater than 0") // panic if the capacity is less than or equal to 0
	}

	logger.Logger.Info().
		Str("name", "MemoryFS").
		Str("capacity", units.HumanSize(float64(capacity))).
		Msg("init")

	mfs := &MemoryFS{
		files:    make(map[string]*file),
		capacity: capacity,
		mu:       sync.RWMutex{},
		keyLocks: sync.Map{},
		LRU:      newLruList(),
	}

	memoryCapacityBytes.Set(float64(capacity))
	memorySizeBytes.Set(float64(mfs.Size()))

	return mfs
}

func (m *MemoryFS) Capacity() int64 {
	return m.capacity
}

func (m *MemoryFS) Name() string {
	return "MemoryFS"
}

func (m *MemoryFS) Size() int64 {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.size
}

func (m *MemoryFS) getKeyLock(key string) *sync.RWMutex {
	mu, _ := m.keyLocks.LoadOrStore(key, &sync.RWMutex{})
	return mu.(*sync.RWMutex)
}

func (m *MemoryFS) Create(key string, size int64) (io.WriteCloser, error) {
	m.mu.RLock()
	if m.capacity > 0 {
		if m.size+size > m.capacity {
			m.mu.RUnlock()
			return nil, vfserror.ErrDiskFull
		}
	}
	m.mu.RUnlock()

	keyMu := m.getKeyLock(key)
	keyMu.Lock()
	defer keyMu.Unlock()

	buf := &bytes.Buffer{}

	return &memWriteCloser{
		Writer: buf,
		onClose: func() error {
			data := buf.Bytes()
			m.mu.Lock()
			if f, exists := m.files[key]; exists {
				m.size -= int64(len(f.data))
				m.LRU.Remove(key)
			}
			fi := vfs.NewFileInfo(key, int64(len(data)), time.Now())
			m.files[key] = &file{
				fileinfo: fi,
				data:     data,
			}
			m.LRU.Add(key, fi)
			m.size += int64(len(data))
			m.mu.Unlock()

			memoryWriteBytes.Add(float64(len(data)))
			memorySizeBytes.Set(float64(m.Size()))

			return nil
		},
	}, nil
}

type memWriteCloser struct {
	io.Writer
	onClose func() error
}

func (wc *memWriteCloser) Close() error {
	return wc.onClose()
}

func (m *MemoryFS) Delete(key string) error {
	keyMu := m.getKeyLock(key)
	keyMu.Lock()
	defer keyMu.Unlock()

	m.mu.Lock()
	f, exists := m.files[key]
	if !exists {
		m.mu.Unlock()
		return vfserror.ErrNotFound
	}
	m.size -= int64(len(f.data))
	m.LRU.Remove(key)
	delete(m.files, key)
	m.mu.Unlock()

	memorySizeBytes.Set(float64(m.Size()))

	return nil
}

func (m *MemoryFS) Open(key string) (io.ReadCloser, error) {
	keyMu := m.getKeyLock(key)
	keyMu.RLock()
	defer keyMu.RUnlock()

	m.mu.Lock()
	f, exists := m.files[key]
	if !exists {
		m.mu.Unlock()
		return nil, vfserror.ErrNotFound
	}
	f.fileinfo.ATime = time.Now()
	m.LRU.MoveToFront(key)
	dataCopy := make([]byte, len(f.data))
	copy(dataCopy, f.data)
	m.mu.Unlock()

	memoryReadBytes.Add(float64(len(dataCopy)))
	memorySizeBytes.Set(float64(m.Size()))

	return io.NopCloser(bytes.NewReader(dataCopy)), nil
}

func (m *MemoryFS) Stat(key string) (*vfs.FileInfo, error) {
	keyMu := m.getKeyLock(key)
	keyMu.RLock()
	defer keyMu.RUnlock()

	m.mu.RLock()
	defer m.mu.RUnlock()

	f, ok := m.files[key]
	if !ok {
		return nil, vfserror.ErrNotFound
	}

	return f.fileinfo, nil
}

func (m *MemoryFS) StatAll() []*vfs.FileInfo {
	m.mu.RLock()
	defer m.mu.RUnlock()

	// hard copy the file info to prevent modification of the original file info or the other way around
	files := make([]*vfs.FileInfo, 0, len(m.files))
	for _, v := range m.files {
		fi := *v.fileinfo
		files = append(files, &fi)
	}

	return files
}