gopeed/pkg/download/storage.go

package download

import (
	"encoding/json"
	"go.etcd.io/bbolt"
	"os"
	"path/filepath"
	"reflect"
	"sync"
)

type Storage interface {
	Setup(buckets []string) error
	Put(bucket string, key string, v any) error
	Get(bucket string, key string, v any) (bool, error)
	List(bucket string, v any) error
	Pop(bucket string, key string, v any) error
	Delete(bucket string, key string) error

	Close() error
	Clear() error
}

func changeValue(p any, v any) {
	if v == nil {
		return
	}
	rp := reflect.ValueOf(p)
	rv := reflect.ValueOf(v)
	if rv.Kind() == reflect.Slice {
		if rv.Len() == 0 {
			return
		}
		// get underlying type
		tp := reflect.TypeOf(p).Elem().Elem()
		for i := 0; i < rv.Len(); i++ {
			// convert to underlying type
			vv := rv.Index(i).Elem().Convert(tp)
			rp.Elem().Set(reflect.Append(rp.Elem(), vv))
		}
	} else if rv.Kind() == reflect.Ptr {
		rp.Elem().Set(rv.Elem())
	} else {
		rp.Elem().Set(rv)
	}
}

var memData = make(map[string]map[string]any)

type MemStorage struct {
	lock *sync.RWMutex
}

func NewMemStorage() *MemStorage {
	return &MemStorage{
		lock: &sync.RWMutex{},
	}
}

func (n *MemStorage) Setup(buckets []string) error {
	for _, bucket := range buckets {
		if _, ok := memData[bucket]; !ok {
			memData[bucket] = make(map[string]any)
		}
	}
	return nil
}

func (n *MemStorage) Put(bucket string, key string, v any) error {
	n.lock.Lock()
	defer n.lock.Unlock()
	memData[bucket][key] = v
	return nil
}

func (n *MemStorage) Get(bucket string, key string, v any) (bool, error) {
	n.lock.RLock()
	defer n.lock.RUnlock()
	if dv, ok := memData[bucket][key]; ok {
		changeValue(v, dv)
		return true, nil
	}
	return false, nil
}

func (n *MemStorage) List(bucket string, v any) error {
	n.lock.RLock()
	defer n.lock.RUnlock()
	data := memData[bucket]
	list := make([]any, 0)
	for _, v := range data {
		list = append(list, v)
	}
	changeValue(v, list)
	return nil
}

func (n *MemStorage) Pop(bucket string, key string, v any) error {
	n.lock.Lock()
	defer n.lock.Unlock()
	data := memData[bucket]
	changeValue(v, data[key])
	delete(data, key)
	return nil
}

func (n *MemStorage) Delete(bucket string, key string) error {
	n.lock.Lock()
	defer n.lock.Unlock()
	delete(memData[bucket], key)
	return nil
}

func (n *MemStorage) Close() error {
	return nil
}

func (n *MemStorage) Clear() error {
	n.lock.Lock()
	defer n.lock.Unlock()
	memData = make(map[string]map[string]any)
	return nil
}

const (
	dbFile = "gopeed.db"
)

type BoltStorage struct {
	db   *bbolt.DB
	path string
}

func NewBoltStorage(dir string) *BoltStorage {
	if err := os.MkdirAll(dir, 0755); err != nil {
		panic(err)
	}
	path := filepath.Join(dir, dbFile)
	db, err := bbolt.Open(path, 0600, nil)
	if err != nil {
		panic(err)
	}
	return &BoltStorage{
		db:   db,
		path: path,
	}
}

func (b *BoltStorage) Setup(buckets []string) error {
	return b.db.Update(func(tx *bbolt.Tx) error {
		for _, bucket := range buckets {
			_, err := tx.CreateBucketIfNotExists([]byte(bucket))
			if err != nil {
				return err
			}
		}
		return nil
	})
}

func (b *BoltStorage) Put(bucket string, key string, v any) error {
	buf, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return b.db.Update(func(tx *bbolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		return b.Put([]byte(key), buf)
	})
}

func (b *BoltStorage) Get(bucket string, key string, v any) (bool, error) {
	var data []byte
	err := b.db.View(func(tx *bbolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		data = b.Get([]byte(key))
		return nil
	})
	if err != nil {
		return false, err
	}
	if data == nil {
		return false, nil
	}
	if err := json.Unmarshal(data, v); err != nil {
		return false, err
	}
	return true, nil
}

func (b *BoltStorage) List(bucket string, v any) error {
	list := make([]any, 0)
	tv := reflect.TypeOf(v).Elem().Elem()
	if err := b.db.View(func(tx *bbolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if err := b.ForEach(func(k, v []byte) error {
			data := reflect.New(tv.Elem()).Interface()
			if err := json.Unmarshal(v, &data); err != nil {
				return err
			}
			list = append(list, data)
			return nil
		}); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return err
	}
	changeValue(v, list)
	return nil
}

func (b *BoltStorage) Pop(bucket string, key string, v any) error {
	var data []byte
	err := b.db.Update(func(tx *bbolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		kb := []byte(key)
		data = b.Get(kb)
		return b.Delete(kb)
	})
	if err != nil {
		return err
	}
	if len(data) == 0 {
		return nil
	}
	return json.Unmarshal(data, v)
}

func (b *BoltStorage) Delete(bucket string, key string) error {
	return b.db.Update(func(tx *bbolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		return b.Delete([]byte(key))
	})
}

func (b *BoltStorage) Close() error {
	return b.db.Close()
}

func (b *BoltStorage) Clear() error {
	if err := b.Close(); err != nil {
		return err
	}
	if err := os.Remove(b.path); err != nil {
		return err
	}
	return nil
}