143 lines
2.7 KiB
Go
143 lines
2.7 KiB
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()
|
|
}
|
|
}
|