comic-hasher/ch/CHDB_bolt.go

package ch

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"slices"
	"time"

	"gitea.narnian.us/lordwelch/comic-hasher/log"

	bolt "go.etcd.io/bbolt"
	"golang.org/x/exp/constraints"
)

type Version string

type Bolt struct {
	Path string
	DB   *bolt.DB
}

func (b Bolt) Close() error {
	return b.DB.Close()
}

// Compact must not be run while the database is actively being used
func (b *Bolt) Compact() error {
	var needsCompaction bool
	b.DB.View(func(tx *bolt.Tx) error {
		var s bolt.BucketStats
		var count int
		if err := tx.ForEach(func(name []byte, b *bolt.Bucket) error {
			if bytes.HasPrefix(name, []byte{}) {
				s.Add(b.Stats())
				count += 1
			}
			return nil
		}); err != nil {
			return err
		}
		needsCompaction = (s.BranchAlloc > 0 && float32(s.BranchInuse)/float32(s.BranchAlloc) < 0.7) || (s.LeafAlloc > 0 && float32(s.LeafInuse)/float32(s.LeafAlloc) < 0.7)
		return nil
	})
	if !needsCompaction {
		return nil
	}
	b.Close()

	db, err := bolt.Open(b.Path, 0o644, &bolt.Options{FreelistType: bolt.FreelistArrayType, Timeout: time.Millisecond * 100, ReadOnly: true})
	if err != nil {
		return fmt.Errorf("failed to open database: %w", err)
	}
	defer db.Close()

	var compactDB *bolt.DB
	compactCHDBPath := b.Path + "compacted.bolt"
	compactDB, err = bolt.Open(compactCHDBPath, 0o644, &bolt.Options{FreelistType: bolt.FreelistArrayType, Timeout: time.Millisecond * 100, NoSync: true})
	if err != nil {
		return fmt.Errorf("failed to open database: %w", err)
	}
	defer compactDB.Close()
	err = bolt.Compact(compactDB, db, 65536)
	if err != nil {
		return fmt.Errorf("failed to compact the database: %w", err)
	}
	err = compactDB.Sync()
	if err != nil {
		return err
	}
	compactDB.Close()
	db.Close()
	err = os.Rename(compactCHDBPath, b.Path)
	if err != nil {
		return err
	}
	b.DB, err = bolt.Open(b.Path, 0o644, &bolt.Options{FreelistType: bolt.FreelistArrayType, Timeout: time.Millisecond * 100})
	if err != nil {
		return err
	}
	return nil
}

type CHDBBolt struct {
	Bolt
	BasePath       string
	deleteExisting bool
	validPaths     map[string]bool
}

const CHDBVERSION = "1.0.0"

func OpenCHDBBolt(path string, deleteExisting bool) (CHDBBolt, error) {
	path, _ = filepath.Abs(path)
	chDBPath := filepath.Join(path, "chdb.bolt")
	bolt, err := OpenBolt(chDBPath, CHDBVERSION, []string{"paths", "bad_urls", "source"})
	if err != nil {
		return CHDBBolt{Bolt: bolt, BasePath: path, deleteExisting: deleteExisting, validPaths: make(map[string]bool)}, fmt.Errorf("failed to open database: %w", err)
	}
	return CHDBBolt{Bolt: bolt, BasePath: path, deleteExisting: deleteExisting, validPaths: make(map[string]bool)}, nil
}

func OpenBolt(path string, version Version, buckets []string) (Bolt, error) {
	path, _ = filepath.Abs(path)
	db, err := bolt.Open(path, 0o644, &bolt.Options{FreelistType: bolt.FreelistArrayType, Timeout: time.Second * 5})
	if err != nil {
		return Bolt{path, db}, fmt.Errorf("failed to open database: %w", err)
	}
	db, err = BoltTruncate(db, version)
	if err != nil {
		return Bolt{}, fmt.Errorf("failed to truncate database")
	}
	err = db.Update(func(tx *bolt.Tx) error {
		for _, bucket := range buckets {
			_, err = tx.CreateBucketIfNotExists([]byte(bucket))
			if err != nil {
				return fmt.Errorf("failed to create bucket %v: %w", bucket, err)
			}
		}

		var versionBucket *bolt.Bucket
		versionBucket, err = tx.CreateBucketIfNotExists([]byte("version"))
		if err != nil {
			return fmt.Errorf("failed to create bucket %v: %w", "paths", err)
		}
		err = versionBucket.Put([]byte("version"), []byte(version))
		if err != nil {
			return fmt.Errorf("failed to set Bolt version: %w", err)
		}
		return nil
	})
	if err != nil {
		db.Close()
		return Bolt{path, db}, fmt.Errorf("failed to init database: %w", err)
	}

	return Bolt{path, db}, nil
}

func BoltTruncate(db *bolt.DB, version Version) (*bolt.DB, error) {
	doTruncate := false
	err := db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("version"))
		if bucket == nil {
			doTruncate = true
			return nil
		}
		if !bytes.Equal(bucket.Get([]byte("version")), []byte(version)) {
			doTruncate = true
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	if !doTruncate {
		return db, nil
	}

	path := db.Path()
	db.Close()
	CHDBLogger.Println(log.Verbose, "Creating new Bolt database")
	err = os.Truncate(path, 0)
	if err != nil {
		return nil, err
	}
	return bolt.Open(path, 0o644, nil)
}

// Compact must not be run while the database is actively being used
func (c CHDBBolt) Compact() error {
	return c.Bolt.Compact()
}

func (c CHDBBolt) Import(paths []string, badURLs []string) {
	slices.Sort(paths)
	slices.Sort(badURLs)
	c.DB.Update(func(tx *bolt.Tx) error {
		p := tx.Bucket([]byte("paths"))
		b := tx.Bucket([]byte("bad_urls"))

		for _, path := range paths {
			p.Put([]byte(path), []byte{})
		}
		for _, url := range badURLs {
			b.Put([]byte(url), []byte{})
		}
		return nil
	})
}

func (c CHDBBolt) ValidatePath(path string) bool {
	firstDir, _ := log.Partition(path, "/")
	if c.validPaths[firstDir] {
		return true
	}
	_, err := os.Stat(filepath.Join(c.BasePath, firstDir))
	c.validPaths[firstDir] = err == nil
	return c.validPaths[firstDir]
}

func (c CHDBBolt) Dump() (paths []string, badURLs []string) {
	c.DB.View(func(tx *bolt.Tx) error {
		p := tx.Bucket([]byte("paths"))
		b := tx.Bucket([]byte("bad_urls"))
		paths = make([]string, 0, p.Inspect().KeyN)
		badURLs = make([]string, 0, b.Inspect().KeyN)
		b.ForEach(func(k, v []byte) error {
			badURLs = append(badURLs, string(k)+"")
			return nil
		})
		p.ForEach(func(k, v []byte) error {
			paths = append(paths, string(k)+"")
			return nil
		})
		return nil
	})
	return paths, badURLs
}

func (c CHDBBolt) PathHashed(path string) bool {
	if filepath.IsAbs(path) {
		panic(fmt.Sprintln("Absolute path given to chdb:", path))
	}
	if !c.ValidatePath(path) {
		panic(fmt.Sprintln("Invalid path", path))
	}

	tx, err := c.DB.Begin(false)
	if err != nil {
		return false
	}
	defer tx.Rollback()
	b := tx.Bucket([]byte("paths"))
	dbRes := b.Get([]byte(path))
	if dbRes != nil {
		if c.deleteExisting {
			err = os.Remove(filepath.Join(c.BasePath, path))
			if err != nil && !errors.Is(err, os.ErrNotExist) {
				CHDBLogger.Println(log.Required, "Failed to delete hashed image:", err)
			}
		}
		return true
	}

	return false
}

func IntToBolt[T constraints.Integer](v T) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, uint64(v))
	return b
}

func (c CHDBBolt) PathDownloaded(path string) bool {
	if filepath.IsAbs(path) {
		panic(fmt.Sprintln("Absolute path given to chdb:", path))
	}
	if !c.ValidatePath(path) {
		panic(fmt.Sprintln("Invalid path", path))
	}

	tx, err := c.DB.Begin(false)
	if err != nil {
		return false
	}
	defer tx.Rollback()
	b := tx.Bucket([]byte("paths"))
	dbRes := b.Get([]byte(path))
	if dbRes == nil {
		_, err := os.Stat(filepath.Join(c.BasePath, path))
		// We only care if the error is not ErrNotExist
		// If it's something else we should get an error on the next operation
		// If it is nil then it still is not ErrNotExist which means we return true
		return !errors.Is(err, fs.ErrNotExist)
	}
	return true
}

func (c CHDBBolt) AddPath(path string) {
	if filepath.IsAbs(path) {
		panic(fmt.Sprintln("Absolute path given to chdb:", path))
	}
	if !c.ValidatePath(path) {
		panic(fmt.Sprintln("Invalid path", path))
	}

	tx, err := c.DB.Begin(true)
	if err != nil {
		CHDBLogger.Log(log.Fatal, "failed to open transaction: %v", err)
	}
	defer tx.Rollback()
	b := tx.Bucket([]byte("paths"))

	err = b.Put([]byte(path), []byte{})
	if err != nil {
		CHDBLogger.Log(log.Fatal, "failed to insert %v (%v) into paths: %v", path, path, err)
	}
	tx.Commit()
	if c.deleteExisting {
		err = os.Remove(filepath.Join(c.BasePath, path))
		if err != nil && !errors.Is(err, os.ErrNotExist) {
			CHDBLogger.Println(log.Required, "Failed to delete hashed image:", err)
		}
		_ = RmdirP(filepath.Dir(path))
	}
}

func (c CHDBBolt) CheckURL(url string) bool {
	tx, err := c.DB.Begin(true)
	if err != nil {
		return false
	}
	defer tx.Rollback()
	b := tx.Bucket([]byte("bad_urls"))
	return b.Get([]byte(url)) != nil
}

func (c CHDBBolt) AddURL(url string) {
	tx, err := c.DB.Begin(true)
	if err != nil {
		CHDBLogger.Log(log.Fatal, "Failed to open transaction: %v", err)
	}
	defer tx.Rollback()
	b := tx.Bucket([]byte("bad_urls"))

	err = b.Put([]byte(url), []byte{})
	if err != nil {
		CHDBLogger.Log(log.Fatal, "failed to insert %v into bad_urls: %v", url, err) // TODO: this should probably be fatal
	}
	tx.Commit()
}

func (c CHDBBolt) Close() error {
	return c.DB.Close()
}

var _ CHDB = CHDBBolt{}