rebound/stat/store.go

package stat

import (
	"encoding/json"
	"os"
	"path/filepath"
	"sync"

	"log/slog"

	"github.com/pkg/errors"
)

type Store struct {
	data         sync.Map
	loadSaveLock sync.Mutex
}

func (s *Store) Load(path string) error {
	s.loadSaveLock.Lock()
	defer s.loadSaveLock.Unlock()

	file, err := os.OpenFile(path, os.O_RDONLY, os.ModePerm)
	if err != nil {
		return errors.WithStack(err)
	}

	decoder := json.NewDecoder(file)
	data := map[string]any{}

	if err := decoder.Decode(&data); err != nil {
		return errors.WithStack(err)
	}

	s.data.Range(func(key, value any) bool {
		s.data.Delete(key)
		return true
	})

	for k, v := range data {
		s.data.Store(k, v)
	}

	return nil
}

func (s *Store) Save(path string) error {
	s.loadSaveLock.Lock()
	defer s.loadSaveLock.Unlock()

	data, err := s.Snapshot()
	if err != nil {
		return errors.WithStack(err)
	}

	dir := filepath.Dir(path)
	filename := filepath.Base(path)

	temp, err := os.CreateTemp(dir, filename+".new*")
	if err != nil {
		return errors.WithStack(err)
	}

	defer func() {
		if err := os.Remove(temp.Name()); err != nil && !errors.Is(err, os.ErrNotExist) {
			slog.Error("could not remove temporary file",
				slog.String("file", temp.Name()),
				slog.Any("error", errors.WithStack(err)),
			)
		}
	}()

	encoder := json.NewEncoder(temp)
	if err := encoder.Encode(data); err != nil {
		return errors.WithStack(err)
	}

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

	return nil
}

func (s *Store) Snapshot() (map[string]float64, error) {
	data := map[string]float64{}

	var err error
	s.data.Range(func(rawKey, rawValue any) bool {
		key, ok := rawKey.(string)
		if !ok {
			err = errors.Errorf("unexpected stat key of '%v'", rawKey)

			return false
		}

		value, ok := rawValue.(float64)
		if !ok {
			err = errors.Errorf("unexpected stat value of '%v'", rawValue)

			return false
		}

		data[key] = value

		return true
	})
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return data, nil
}

func (s *Store) Add(name string, added float64, defaultValue float64) float64 {
	for {
		value := s.Get(name, defaultValue)
		if value == defaultValue {
			s.data.Store(name, defaultValue)
		}

		sum := value + added
		if s.data.CompareAndSwap(name, value, value+added) {
			return sum
		}
	}
}

func (s *Store) Set(name string, value float64) float64 {
	s.data.Store(name, value)

	return value
}

func (s *Store) Get(name string, defaultValue float64) float64 {
	rawValue, ok := s.data.Load(name)
	if !ok {
		return defaultValue
	}

	value, ok := rawValue.(float64)
	if !ok {
		return defaultValue
	}

	return value
}

func NewStore() *Store {
	return &Store{
		data:         sync.Map{},
		loadSaveLock: sync.Mutex{},
	}
}