star/gdu/pkg/analyze/storage.go

package analyze

import (
	"bytes"
	"encoding/gob"
	"path/filepath"
	"sync"

	"b612.me/apps/b612/gdu/pkg/fs"
	"github.com/dgraph-io/badger/v3"
	"github.com/pkg/errors"
)

func init() {
	gob.RegisterName("analyze.StoredDir", &StoredDir{})
	gob.RegisterName("analyze.Dir", &Dir{})
	gob.RegisterName("analyze.File", &File{})
	gob.RegisterName("analyze.ParentDir", &ParentDir{})
}

// DefaultStorage is a default instance of badger storage
var DefaultStorage *Storage

// Storage represents a badger storage
type Storage struct {
	db          *badger.DB
	storagePath string
	topDir      string
	m           sync.RWMutex
	counter     int
	counterM    sync.Mutex
}

// NewStorage returns new instance of badger storage
func NewStorage(storagePath, topDir string) *Storage {
	st := &Storage{
		storagePath: storagePath,
		topDir:      topDir,
	}
	DefaultStorage = st
	return st
}

// GetTopDir returns top directory
func (s *Storage) GetTopDir() string {
	return s.topDir
}

// IsOpen returns true if badger DB is open
func (s *Storage) IsOpen() bool {
	s.m.RLock()
	defer s.m.RUnlock()
	return s.db != nil
}

// Open opens badger DB
func (s *Storage) Open() func() {
	options := badger.DefaultOptions(s.storagePath)
	options.Logger = nil
	db, err := badger.Open(options)
	if err != nil {
		panic(err)
	}
	s.db = db

	return func() {
		s.db.Close()
		s.db = nil
	}
}

// StoreDir saves item info into badger DB
func (s *Storage) StoreDir(dir fs.Item) error {
	s.checkCount()
	s.m.RLock()
	defer s.m.RUnlock()

	return s.db.Update(func(txn *badger.Txn) error {
		b := &bytes.Buffer{}
		enc := gob.NewEncoder(b)
		err := enc.Encode(dir)
		if err != nil {
			return errors.Wrap(err, "encoding dir value")
		}

		return txn.Set([]byte(dir.GetPath()), b.Bytes())
	})
}

// LoadDir saves item info into badger DB
func (s *Storage) LoadDir(dir fs.Item) error {
	s.checkCount()
	s.m.RLock()
	defer s.m.RUnlock()

	return s.db.View(func(txn *badger.Txn) error {
		path := dir.GetPath()
		item, err := txn.Get([]byte(path))
		if err != nil {
			return errors.Wrap(err, "reading stored value for path: "+path)
		}
		return item.Value(func(val []byte) error {
			b := bytes.NewBuffer(val)
			dec := gob.NewDecoder(b)
			return dec.Decode(dir)
		})
	})
}

// GetDirForPath returns Dir for given path
func (s *Storage) GetDirForPath(path string) (fs.Item, error) {
	dirPath := filepath.Dir(path)
	name := filepath.Base(path)
	dir := &StoredDir{
		&Dir{
			File: &File{
				Name: name,
			},
			BasePath: dirPath,
		},
		nil,
		sync.Mutex{},
	}
	err := s.LoadDir(dir)
	if err != nil {
		return nil, err
	}
	return dir, nil
}

func (s *Storage) checkCount() {
	s.counterM.Lock()
	defer s.counterM.Unlock()
	s.counter++
	if s.counter >= 10000 {
		s.m.Lock()
		defer s.m.Unlock()
		s.counter = 0
		s.db.Close()
		s.Open()
	}
}