switched to rosedb, it actually works! stopped using goroutines for file indexing since its slower

1 files changed, 81 insertions, 60 deletions

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
+	}
+
+	info, err := decodeDbInfo(value)
+	if err != nil {
 		return nil, err
 	}
-	return info, nil
+	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
 }