cost-model/core/pkg/storage/memorystorage.go

package storage

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

	"github.com/opencost/opencost/core/pkg/log"
	"github.com/opencost/opencost/core/pkg/storage/memfile"
)

// MemoryStorage is a thread-safe in-memory file system storage implementation. It can be used for testing storage.Storage dependents
// or to serve as a lightweight storage implementation within a production system.
type MemoryStorage struct {
	lock        sync.Mutex
	directPaths map[string]*memfile.MemoryFile
	fileTree    *memfile.MemoryDirectory
}

// NewMemoryStorage creates a new in-memory file system storage implementation.
func NewMemoryStorage() *MemoryStorage {
	return &MemoryStorage{
		directPaths: make(map[string]*memfile.MemoryFile),
		fileTree:    memfile.NewMemoryDirectory(""),
	}
}

// String returns the storage type as a string for logging purposes.
func (ms *MemoryStorage) String() string {
	return string(ms.StorageType())
}

// StorageType returns a string identifier for the type of storage used by the implementation.
func (ms *MemoryStorage) StorageType() StorageType {
	return StorageTypeMemory
}

// FullPath returns the storage working path combined with the path provided
func (ms *MemoryStorage) FullPath(path string) string {
	return path
}

// Stat returns the StorageStats for the specific path.
func (ms *MemoryStorage) Stat(path string) (*StorageInfo, error) {
	ms.lock.Lock()
	defer ms.lock.Unlock()

	path = filepath.Clean(path)
	if file, ok := ms.directPaths[path]; ok {
		return &StorageInfo{
			Name:    file.Name,
			Size:    file.Size(),
			ModTime: file.ModTime,
		}, nil
	}

	return nil, DoesNotExistError
}

// Read uses the relative path of the storage combined with the provided path to
// read the contents.
func (ms *MemoryStorage) Read(path string) ([]byte, error) {
	ms.lock.Lock()
	defer ms.lock.Unlock()

	path = filepath.Clean(path)

	if file, ok := ms.directPaths[path]; ok {
		return file.Contents, nil
	}

	return nil, DoesNotExistError
}

// ReadStream returns a streaming reader for the specified in-memory object.
func (ms *MemoryStorage) ReadStream(path string) (io.ReadCloser, error) {
	ms.lock.Lock()
	defer ms.lock.Unlock()

	path = filepath.Clean(path)

	if file, ok := ms.directPaths[path]; ok {
		data := append([]byte(nil), file.Contents...)
		return io.NopCloser(bytes.NewReader(data)), nil
	}

	return nil, DoesNotExistError
}

// ReadToLocalFile writes the specified object at path to destPath on the local file system.
func (ms *MemoryStorage) ReadToLocalFile(path, destPath string) error {
	ms.lock.Lock()
	path = filepath.Clean(path)

	file, ok := ms.directPaths[path]
	if !ok {
		ms.lock.Unlock()
		return DoesNotExistError
	}

	// Copy the contents so we can release the lock before doing potentially slow disk IO.
	data := append([]byte(nil), file.Contents...)
	ms.lock.Unlock()

	dir := filepath.Dir(destPath)
	if err := os.MkdirAll(dir, os.ModePerm); err != nil {
		return fmt.Errorf("MemoryStorage: ReadToLocalFile: creating destination directory: %w", err)
	}
	if err := os.WriteFile(destPath, data, 0600); err != nil {
		return fmt.Errorf("MemoryStorage: ReadToLocalFile: writing destination file: %w", err)
	}

	return nil
}

// Write uses the relative path of the storage combined with the provided path
// to write a new file or overwrite an existing file.
func (ms *MemoryStorage) Write(path string, data []byte) error {
	ms.lock.Lock()
	defer ms.lock.Unlock()

	paths, pFile := memfile.Split(path)

	f := memfile.NewMemoryFile(pFile, data)
	currentDir := memfile.CreateSubdirectory(ms.fileTree, paths)

	currentDir.AddFile(f)
	ms.directPaths[path] = f
	return nil
}

// WriteStream creates a new relative path and returns the io.WriteCloser that can be used to
// write into the storage path. Close() blocks until all data has been committed to storage.
func (ms *MemoryStorage) WriteStream(path string) (io.WriteCloser, error) {
	r, w := io.Pipe()
	var wg sync.WaitGroup

	wg.Go(func() {
		data, err := io.ReadAll(r)
		if err != nil {
			r.CloseWithError(err)
			return
		}

		ms.lock.Lock()
		defer ms.lock.Unlock()

		paths, pFile := memfile.Split(path)

		f := memfile.NewMemoryFile(pFile, data)
		currentDir := memfile.CreateSubdirectory(ms.fileTree, paths)

		currentDir.AddFile(f)
		ms.directPaths[path] = f
	})

	return &memWriteCloser{
		PipeWriter: w,
		wg:         &wg,
	}, nil
}

// memWriteCloser wraps *io.PipeWriter so that Close() blocks until the drain goroutine
// has finished committing data to the in-memory store.
type memWriteCloser struct {
	*io.PipeWriter
	wg *sync.WaitGroup
}

func (m *memWriteCloser) Close() error {
	err := m.PipeWriter.Close()
	m.wg.Wait()
	return err
}

// Remove uses the relative path of the storage combined with the provided path to
// remove a file from storage permanently.
func (ms *MemoryStorage) Remove(path string) error {
	ms.lock.Lock()
	defer ms.lock.Unlock()

	path = filepath.Clean(path)
	paths, pFile := memfile.Split(path)

	currentDir, err := memfile.FindSubdirectory(ms.fileTree, paths)
	if err != nil {
		return fmt.Errorf("file not found: %s - %w", path, DoesNotExistError)
	}

	currentDir.RemoveFile(pFile)

	delete(ms.directPaths, path)
	return nil
}

// Exists uses the relative path of the storage combined with the provided path to
// determine if the file exists.
func (ms *MemoryStorage) Exists(path string) (bool, error) {
	ms.lock.Lock()
	defer ms.lock.Unlock()

	path = filepath.Clean(path)

	_, ok := ms.directPaths[path]
	return ok, nil
}

// List uses the relative path of the storage combined with the provided path to return
// storage information for the files.
func (ms *MemoryStorage) List(path string) ([]*StorageInfo, error) {
	ms.lock.Lock()
	defer ms.lock.Unlock()

	paths := memfile.SplitPaths(path)
	currentDir, err := memfile.FindSubdirectory(ms.fileTree, paths)
	if err != nil {
		// contract for bucket storages returns an empty list in this case
		// so just log a warning, and return an empty list
		log.Warnf("failed to resolve path: %s - %s", path, err)
		return []*StorageInfo{}, nil
	}

	storageInfos := make([]*StorageInfo, 0, currentDir.FileCount())
	for f := range currentDir.Files() {
		storageInfos = append(storageInfos, &StorageInfo{
			Name:    f.Name,
			Size:    f.Size(),
			ModTime: f.ModTime,
		})
	}

	return storageInfos, nil
}

// ListDirectories uses the relative path of the storage combined with the provided path
// to return storage information for only directories contained along the path. This
// functions as List, but returns storage information for only directories.
func (ms *MemoryStorage) ListDirectories(path string) ([]*StorageInfo, error) {
	ms.lock.Lock()
	defer ms.lock.Unlock()

	paths := memfile.SplitPaths(path)
	currentDir, err := memfile.FindSubdirectory(ms.fileTree, paths)
	if err != nil {
		// contract for bucket storages returns an empty list in this case
		// so just log a warning, and return an empty list
		log.Warnf("failed to resolve path: %s - %s", path, err)
		return []*StorageInfo{}, nil
	}

	storageInfos := make([]*StorageInfo, 0, currentDir.DirCount())
	for d := range currentDir.Directories() {
		storageInfos = append(storageInfos, &StorageInfo{
			Name:    filepath.Join(append(paths, d.Name)...) + "/",
			Size:    d.Size(),
			ModTime: d.ModTime,
		})
	}

	return storageInfos, nil
}