From b1a2ab25856ec14b042da9fec49c453c6dad95e5 Mon Sep 17 00:00:00 2001 From: jakka Date: Sun, 13 Apr 2025 13:59:04 +0300 Subject: switched to rosedb, it actually works! stopped using goroutines for file indexing since its slower --- files/database.go | 143 +++++++++++++++++++++++++++++++----------------------- 1 file changed, 82 insertions(+), 61 deletions(-) (limited to 'files/database.go') diff --git a/files/database.go b/files/database.go index 5cb8e98..33c2636 100644 --- a/files/database.go +++ b/files/database.go @@ -3,46 +3,48 @@ package files import ( "context" "encoding/json" + "errors" + "log" + "os" + "strings" - badger "github.com/dgraph-io/badger/v4" + "github.com/rosedblabs/rosedb/v2" progressbar "github.com/schollz/progressbar/v3" "golang.org/x/sync/errgroup" ) -func decodeDbInfo(item *badger.Item) (*FileInfo, error) { +func decodeDbInfo(value []byte) (*FileInfo, error) { var info FileInfo - if err := item.Value(func(val []byte) error { - if err := json.Unmarshal(val, &info); err != nil { - return err - } - return nil - }); err != nil { + if err := json.Unmarshal(value, &info); err != nil { return nil, err } return &info, nil } -func getInfoFromDb(database *badger.DB, hashsum []byte) (*FileInfo, error) { - var info *FileInfo - if err := database.View(func(txn *badger.Txn) error { - item, err := txn.Get(hashsum) - if err != nil { - return err - } +func getInfoFromDb(batch *rosedb.Batch, hashsum []byte) (*FileInfo, error) { + exist, err := batch.Exist(hashsum) + if err != nil { + return nil, err + } - info, err = decodeDbInfo(item) - if err != nil { - return err - } - return nil - }); err != nil { + if !exist { + return nil, rosedb.ErrKeyNotFound + } + + value, err := batch.Get(hashsum) + if err != nil { return nil, err } - return info, nil + + info, err := decodeDbInfo(value) + if err != nil { + return nil, err + } + return info, err } -func (filedata *FileInfo) evaluateFile(database *badger.DB, hashsum []byte, encoder string) error { - info, err := getInfoFromDb(database, hashsum) +func (filedata *FileInfo) evaluateFile(batch *rosedb.Batch, hashsum []byte, encoder string) error { + info, err := getInfoFromDb(batch, hashsum) if err != nil { return err } @@ -59,71 +61,79 @@ func (filedata *FileInfo) evaluateFile(database *badger.DB, hashsum []byte, enco return ReencodeNotNeeded } -func (filedata *FileInfo) updateFile(wb *badger.WriteBatch, hashsum []byte) error { +func (filedata *FileInfo) updateFile(batch *rosedb.Batch, hashsum []byte) error { encoded, err := json.Marshal(filedata) if err != nil { return err } - return wb.Set(hashsum, encoded) + return batch.Put(hashsum, encoded) } -func (filedata *FileInfo) IndexFile(ctx context.Context, hashsum []byte) error { - database := ctx.Value("database").(*badger.DB) +func (filedata *FileInfo) IndexFile(ctx context.Context, hashsum []byte, batch *rosedb.Batch) error { encoder := ctx.Value("encoder").(string) counter := ctx.Value("counter").(*int64) - err := filedata.evaluateFile(database, hashsum, encoder) + err := filedata.evaluateFile(batch, hashsum, encoder) + switch err { case FileMoved, ReencodeNotNeeded: filedata.Process = false - case ReencodeNeeded, badger.ErrKeyNotFound: + case ReencodeNeeded, rosedb.ErrKeyNotFound: *counter += 1 default: return err } - wb := database.NewWriteBatch() - defer wb.Cancel() - if err := filedata.updateFile(wb, hashsum); err != nil { - return err - } - if err := wb.Flush(); err != nil { + if err := filedata.updateFile(batch, hashsum); err != nil { return err } + return nil } func ReencodeFlacs(ctx context.Context) error { - db := ctx.Value("database").(*badger.DB) + db := ctx.Value("database").(*rosedb.DB) bar := progressbar.NewOptions64( *ctx.Value("counter").(*int64), progressbar.OptionSetDescription("Reencoding..."), progressbar.OptionShowCount(), ) - defer bar.Close() - wb := db.NewWriteBatch() - defer wb.Cancel() + batch := db.NewBatch(rosedb.DefaultBatchOptions) wg := new(errgroup.Group) wg.SetLimit(4) - err := db.View(func(txn *badger.Txn) error { - opts := badger.DefaultIteratorOptions - opts.PrefetchSize = 10 - it := txn.NewIterator(opts) - defer it.Close() - - for it.Rewind(); it.Valid(); it.Next() { - select { - case <-ctx.Done(): - return nil - default: - item := it.Item() - key := item.Key() - info, err := decodeDbInfo(item) - if err != nil { - return err + + iterOpts := rosedb.DefaultIteratorOptions + iterOpts.ContinueOnError = true + iter := db.NewIterator(iterOpts) + + for iter.Rewind(); iter.Valid(); iter.Next() { + select { + case <-ctx.Done(): + wg.Wait() + if err := batch.Commit(); err != nil { + return err + } + bar.Exit() + return nil + default: + item := iter.Item() + key := item.Key + info, err := decodeDbInfo(item.Value) + if err != nil { + log.Println(err.Error()) + } + + if !strings.Contains(info.AbsPath, ctx.Value("path").(string)) { + continue + } else if _, err := os.Stat(info.AbsPath); errors.Is(err, os.ErrNotExist) { + if err := batch.Delete(key); err != nil { + log.Println(err.Error()) } + } else if err != nil { + log.Println(err.Error()) + } else { if info.Process { wg.Go(func() error { select { @@ -133,19 +143,30 @@ func ReencodeFlacs(ctx context.Context) error { if err := info.reencodeFile(ctx); err != nil { return err } - if err := info.updateFile(wb, key); err != nil { + if err := batch.Delete(key); err != nil { + return err + } + key, err = getSha256(info.AbsPath) + if err != nil { + return err + } + if err := info.updateFile(batch, key); err != nil { return err } + bar.Add64(1) return nil } }) } - } } - return nil - }) + } + wg.Wait() - return err + if err := batch.Commit(); err != nil { + return err + } + bar.Close() + return nil } -- cgit v1.3.1