129 lines
2.3 KiB
Go
129 lines
2.3 KiB
Go
package internal
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"git.mills.io/prologic/bitcask"
|
|
log "github.com/sirupsen/logrus"
|
|
)
|
|
|
|
const (
|
|
urlsKeyPrefix = "/urls"
|
|
)
|
|
|
|
// BitcaskStore ...
|
|
type BitcaskStore struct {
|
|
db *bitcask.Bitcask
|
|
}
|
|
|
|
func newBitcaskStore(path string) (*BitcaskStore, error) {
|
|
db, err := bitcask.Open(
|
|
path,
|
|
bitcask.WithMaxKeySize(256),
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &BitcaskStore{db: db}, nil
|
|
}
|
|
|
|
// Sync ...
|
|
func (bs *BitcaskStore) Sync() error {
|
|
return bs.db.Sync()
|
|
}
|
|
|
|
// Close ...
|
|
func (bs *BitcaskStore) Close() error {
|
|
log.Info("syncing store ...")
|
|
if err := bs.db.Sync(); err != nil {
|
|
log.WithError(err).Error("error syncing store")
|
|
return err
|
|
}
|
|
|
|
log.Info("closing store ...")
|
|
if err := bs.db.Close(); err != nil {
|
|
log.WithError(err).Error("error closing store")
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// Merge ...
|
|
func (bs *BitcaskStore) Merge() error {
|
|
log.Info("merging store ...")
|
|
if err := bs.db.Merge(); err != nil {
|
|
log.WithError(err).Error("error merging store")
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (bs *BitcaskStore) HasURL(hash string) bool {
|
|
key := []byte(fmt.Sprintf("%s/%s", urlsKeyPrefix, hash))
|
|
return bs.db.Has(key)
|
|
}
|
|
|
|
func (bs *BitcaskStore) DelURL(hash string) error {
|
|
key := []byte(fmt.Sprintf("%s/%s", urlsKeyPrefix, hash))
|
|
return bs.db.Delete(key)
|
|
}
|
|
|
|
func (bs *BitcaskStore) GetURL(hash string) (*URL, error) {
|
|
key := []byte(fmt.Sprintf("%s/%s", urlsKeyPrefix, hash))
|
|
data, err := bs.db.Get(key)
|
|
if err == bitcask.ErrKeyNotFound {
|
|
return nil, ErrURLNotFound
|
|
}
|
|
return LoadURL(data)
|
|
}
|
|
|
|
func (bs *BitcaskStore) SetURL(hash string, url *URL) error {
|
|
data, err := url.Bytes()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
key := []byte(fmt.Sprintf("%s/%s", urlsKeyPrefix, hash))
|
|
if err := bs.db.Put(key, data); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (bs *BitcaskStore) URLCount() int64 {
|
|
var count int64
|
|
|
|
if err := bs.db.Scan([]byte(urlsKeyPrefix), func(_ []byte) error {
|
|
count++
|
|
return nil
|
|
}); err != nil {
|
|
log.WithError(err).Error("error scanning")
|
|
}
|
|
|
|
return count
|
|
}
|
|
|
|
func (bs *BitcaskStore) ForEachURL(f func(url *URL) error) error {
|
|
err := bs.db.Scan([]byte(urlsKeyPrefix), func(key []byte) error {
|
|
data, err := bs.db.Get(key)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
url, err := LoadURL(data)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return f(url)
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|