steamcache2/vfs/disk/disk.go

package disk

import (
	"log"
	"os"
	"path/filepath"
	"s1d3sw1ped/SteamCache2/vfs"
	"s1d3sw1ped/SteamCache2/vfs/vfserror"
	"sync"
	"time"

	"github.com/docker/go-units"
)

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

// DiskFS is a virtual file system that stores files on disk.
type DiskFS struct {
	root string

	info     map[string]*vfs.FileInfo
	capacity int64
	mu       sync.Mutex
	sg       sync.WaitGroup
}

// New creates a new DiskFS.
func new(root string, capacity int64, skipinit bool) *DiskFS {
	dfs := &DiskFS{
		root:     root,
		info:     make(map[string]*vfs.FileInfo),
		capacity: capacity,
		mu:       sync.Mutex{},
		sg:       sync.WaitGroup{},
	}

	os.MkdirAll(dfs.root, 0755)

	if !skipinit {
		dfs.init()
	}

	return dfs
}

func New(root string, capacity int64) *DiskFS {
	return new(root, capacity, false)
}

func NewSkipInit(root string, capacity int64) *DiskFS {
	return new(root, capacity, true)
}

func (d *DiskFS) init() {
	// log.Printf("DiskFS(%s, %s) init", d.root, units.HumanSize(float64(d.capacity)))

	tstart := time.Now()

	d.walk(d.root)
	d.sg.Wait()

	log.Printf("DiskFS(%s, %s) init took %v", d.root, units.HumanSize(float64(d.capacity)), time.Since(tstart))
}

func (d *DiskFS) walk(path string) {
	d.sg.Add(1)
	go func() {
		defer d.sg.Done()
		filepath.Walk(path, func(npath string, info os.FileInfo, err error) error {
			if path == npath {
				return nil
			}

			if err != nil {
				return err
			}

			if info.IsDir() {
				d.walk(npath)
				return filepath.SkipDir
			}

			d.mu.Lock()
			k := npath[len(d.root)+1:]
			d.info[k] = vfs.NewFileInfoFromOS(info, k)
			d.mu.Unlock()

			// log.Printf("DiskFS(%s, %s) init: %s", d.root, units.HumanSize(float64(d.capacity)), npath)
			return nil
		})
	}()
}

func (d *DiskFS) Capacity() int64 {
	return d.capacity
}

func (d *DiskFS) Name() string {
	return "DiskFS"
}

func (d *DiskFS) Size() int64 {
	var size int64
	d.mu.Lock()
	defer d.mu.Unlock()
	for _, v := range d.info {
		size += v.Size()
	}
	return size
}

func (d *DiskFS) Set(key string, src []byte) error {
	if d.capacity > 0 {
		if size := d.Size() + int64(len(src)); size > d.capacity {
			return vfserror.ErrDiskFull
		}
	}

	if _, err := d.Stat(key); err == nil {
		d.Delete(key)
	}

	d.mu.Lock()
	defer d.mu.Unlock()
	os.MkdirAll(filepath.Join(d.root, filepath.Dir(key)), 0755)
	if err := os.WriteFile(filepath.Join(d.root, key), src, 0644); err != nil {
		return err
	}

	fi, err := os.Stat(filepath.Join(d.root, key))
	if err != nil {
		panic(err)
	}

	d.info[key] = vfs.NewFileInfoFromOS(fi, key)

	return nil
}

// Delete deletes the value of key.
func (d *DiskFS) Delete(key string) error {
	_, err := d.Stat(key)
	if err != nil {
		return err
	}

	d.mu.Lock()
	defer d.mu.Unlock()
	delete(d.info, key)
	if err := os.Remove(filepath.Join(d.root, key)); err != nil {
		return err
	}

	return nil
}

// Get gets the value of key and returns it.
func (d *DiskFS) Get(key string) ([]byte, error) {
	_, err := d.Stat(key)
	if err != nil {
		return nil, err
	}

	d.mu.Lock()
	defer d.mu.Unlock()

	data, err := os.ReadFile(filepath.Join(d.root, key))
	if err != nil {
		return nil, err
	}

	return data, nil
}

// Stat returns the FileInfo of key. If key is not found in the cache, it will stat the file on disk. If the file is not found on disk, it will return vfs.ErrNotFound.
func (d *DiskFS) Stat(key string) (*vfs.FileInfo, error) {
	d.mu.Lock()
	fi, ok := d.info[key]
	d.mu.Unlock() // unlock before statting the file

	if !ok {
		fii, err := os.Stat(filepath.Join(d.root, key))
		if err != nil {
			return nil, vfserror.ErrNotFound
		}

		d.mu.Lock()         // relock to update the info map
		defer d.mu.Unlock() // nothing else needs to unlock before returning

		d.info[key] = vfs.NewFileInfoFromOS(fii, key)
		fi = d.info[key]
		// fallthrough to return fi with shiny new info
	}

	return fi, nil
}

func (m *DiskFS) StatAll() []*vfs.FileInfo {
	m.mu.Lock()
	defer m.mu.Unlock()

	// 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.info))
	for _, v := range m.info {
		fi := *v
		files = append(files, &fi)
	}

	return files
}