edge/pkg/storage/driver/cache/lfu/fs/store.go

package fs

import (
	"fmt"
	"io"
	"os"
	"path/filepath"

	"forge.cadoles.com/arcad/edge/pkg/storage/driver/cache/lfu"
	"github.com/pkg/errors"
)

type Store[K comparable, V any] struct {
	baseDir        string
	getPath        GetPathFunc[K]
	marshalValue   MarshalValueFunc[V]
	unmarshalValue UnmarshalValueFunc[V]
}

// Delete implements Store.
func (s *Store[K, V]) Delete(key K) error {
	path, err := s.getEntryPath(key)
	if err != nil {
		return errors.WithStack(err)
	}

	if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) {
		return errors.WithStack(err)
	}

	return nil
}

// Get implements Store.
func (s *Store[K, V]) Get(key K) (V, error) {
	path, err := s.getEntryPath(key)
	if err != nil {
		return *new(V), errors.WithStack(err)
	}

	value, err := s.readValue(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return *new(V), errors.WithStack(lfu.ErrNotFound)
		}

		return *new(V), errors.WithStack(err)
	}

	return value, nil
}

// Set implements Store.
func (s *Store[K, V]) Set(key K, value V) error {
	path, err := s.getEntryPath(key)
	if err != nil {
		return errors.WithStack(err)
	}

	if err := s.writeValue(path, value); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (s *Store[K, V]) Clear() error {
	if err := os.RemoveAll(s.baseDir); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (s *Store[K, V]) getEntryPath(k K) (string, error) {
	path, err := s.getPath(k)
	if err != nil {
		return "", errors.WithStack(err)
	}

	path = append([]string{s.baseDir}, path...)
	return filepath.Join(path...), nil
}

func (s *Store[K, V]) writeValue(path string, value V) error {
	fi, err := os.Stat(path)
	if err == nil && !fi.Mode().IsRegular() {
		return fmt.Errorf("%s already exists and is not a regular file", path)
	}

	dir := filepath.Dir(path)

	if err := os.MkdirAll(dir, 0750); err != nil {
		return errors.WithStack(err)
	}

	f, err := os.CreateTemp(dir, filepath.Base(path)+".tmp")
	if err != nil {
		return errors.WithStack(err)
	}

	tmpName := f.Name()
	defer func() {
		if err != nil {
			f.Close()
			os.Remove(tmpName)
		}
	}()

	reader, err := s.marshalValue(value)
	if err != nil {
		return errors.WithStack(err)
	}

	if _, err := io.Copy(f, reader); err != nil {
		return errors.WithStack(err)
	}

	if err := f.Sync(); err != nil {
		return errors.WithStack(err)
	}

	if err := f.Close(); err != nil {
		return errors.WithStack(err)
	}

	if err := os.Rename(tmpName, path); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (s *Store[K, V]) readValue(path string) (V, error) {
	file, err := os.Open(path)
	if err != nil {
		return *new(V), errors.WithStack(err)
	}

	defer func() {
		if err := file.Close(); err != nil && !errors.Is(err, os.ErrClosed) {
			panic(errors.WithStack(err))
		}
	}()

	value, err := s.unmarshalValue(file)

	if err != nil {
		return *new(V), errors.WithStack(err)
	}

	return value, nil
}

func NewStore[K comparable, V any](baseDir string, funcs ...OptionsFunc[K, V]) *Store[K, V] {
	opts := DefaultOptions[K, V](funcs...)
	return &Store[K, V]{
		baseDir:        baseDir,
		getPath:        opts.GetPath,
		unmarshalValue: opts.UnmarshalValue,
		marshalValue:   opts.MarshalValue,
	}
}

var _ lfu.Store[string, int] = &Store[string, int]{}