any-sync/client/filestorage/badgerfilestore/badgerfileindex.go

package badgerfilestore

import (
	"bytes"
	"errors"
	"github.com/dgraph-io/badger/v3"
	"github.com/ipfs/go-cid"
	"sync"
	"time"
)

const keyIndexPrefix = "files/indexes/"

type Op string

const (
	OpAdd    Op = "add"
	OpDelete Op = "del"
	OpLoad   Op = "load"
)

func NewFileBadgerIndex(db *badger.DB) *FileBadgerIndex {
	return &FileBadgerIndex{
		db:     db,
		workCh: make(chan struct{}, 1),
	}
}

var cidsPool = &sync.Pool{
	New: func() any {
		return &Cids{}
	},
}

type FileBadgerIndex struct {
	db     *badger.DB
	workCh chan struct{}
}

func (i *FileBadgerIndex) Add(cids *Cids) error {
	addTimeBin, _ := time.Now().MarshalBinary()
	defer i.pingWorkCh()
	return i.db.Update(func(txn *badger.Txn) error {
		keys := cids.Keys()
		for _, k := range keys {
			if err := txn.Set(k, addTimeBin); err != nil {
				return err
			}
		}
		return nil
	})
}

func (i *FileBadgerIndex) Done(cids *Cids) (err error) {
	if len(cids.SpaceOps) == 0 {
		return nil
	}
	return i.db.Update(func(txn *badger.Txn) error {
		keys := cids.Keys()
		for _, k := range keys {
			if err := txn.Delete(k); err != nil {
				return err
			}
		}
		return nil
	})
}
func (i *FileBadgerIndex) List(limit int) (cids *Cids, err error) {
	cids = NewCids()
	err = i.db.View(func(txn *badger.Txn) error {
		it := txn.NewIterator(badger.IteratorOptions{
			PrefetchSize:   limit,
			PrefetchValues: false,
			Prefix:         []byte(keyIndexPrefix),
		})
		defer it.Close()
		var l int
		for it.Rewind(); it.Valid(); it.Next() {
			e := cids.AddKey(it.Item().Key())
			if e == nil {
				l++
				if l == limit {
					return nil
				}
			}
		}
		return nil
	})
	return
}

func (i *FileBadgerIndex) Len() (l int, err error) {
	err = i.db.View(func(txn *badger.Txn) error {
		it := txn.NewIterator(badger.IteratorOptions{
			PrefetchSize:   100,
			PrefetchValues: false,
			Prefix:         []byte(keyIndexPrefix),
		})
		defer it.Close()
		var l int
		for it.Rewind(); it.Valid(); it.Next() {
			l++
		}
		return nil
	})
	return
}

func (i *FileBadgerIndex) pingWorkCh() {
	select {
	case i.workCh <- struct{}{}:
	default:
	}
}

func (i *FileBadgerIndex) HasWorkCh() chan struct{} {
	l, err := i.Len()
	if err != nil {
		return i.workCh
	}
	if l > 0 {
		i.pingWorkCh()
	}
	return i.workCh
}

var errInvalidKey = errors.New("invalid key")

var sep = []byte("/")

func parseCIDOp(key []byte) (spaceId string, op Op, k cid.Cid, err error) {
	if len(key) <= len(keyIndexPrefix) {
		err = errInvalidKey
		return
	}
	key = key[len(keyIndexPrefix):]
	fi := bytes.Index(key, sep)
	if fi < 0 {
		err = errInvalidKey
		return
	}
	spaceId = string(key[:fi])
	key = key[fi+1:]
	fi = bytes.Index(key, sep)
	if fi < 0 {
		err = errInvalidKey
		return
	}
	op = Op(key[:fi])
	k, err = cid.Cast(key[fi+1:])
	return
}

func NewCids() *Cids {
	return cidsPool.Get().(*Cids)
}

type SpaceCidOps struct {
	SpaceId string
	Add     []cid.Cid
	Delete  []cid.Cid
	Load    []cid.Cid
}

type Cids struct {
	SpaceOps []SpaceCidOps
	keysBuf  [][]byte
}

func (c *Cids) Add(spaceId string, op Op, k cid.Cid) {
	var spaceIndex = -1
	for i, so := range c.SpaceOps {
		if so.SpaceId == spaceId {
			spaceIndex = i
			break
		}
	}
	if spaceIndex == -1 {
		spaceIndex = len(c.SpaceOps)
		if len(c.SpaceOps) < cap(c.SpaceOps) {
			c.SpaceOps = c.SpaceOps[0 : len(c.SpaceOps)+1]
			c.SpaceOps[spaceIndex].SpaceId = spaceId
		} else {
			c.SpaceOps = append(c.SpaceOps, SpaceCidOps{SpaceId: spaceId})
		}
	}
	switch op {
	case OpAdd:
		c.SpaceOps[spaceIndex].Add = append(c.SpaceOps[spaceIndex].Add, k)
	case OpDelete:
		c.SpaceOps[spaceIndex].Delete = append(c.SpaceOps[spaceIndex].Delete, k)
	case OpLoad:
		c.SpaceOps[spaceIndex].Load = append(c.SpaceOps[spaceIndex].Load, k)
	}
}

func (c *Cids) Keys() [][]byte {
	addKey := func(spaceId string, k cid.Cid, op Op) {
		if len(c.keysBuf) < cap(c.keysBuf) {
			c.keysBuf = c.keysBuf[:len(c.keysBuf)+1]
		} else {
			c.keysBuf = append(c.keysBuf, nil)
		}
		ki := len(c.keysBuf) - 1
		buf := bytes.NewBuffer(c.keysBuf[ki][:0])
		buf.WriteString(keyIndexPrefix)
		buf.WriteString(spaceId)
		buf.WriteString("/")
		buf.WriteString(string(op))
		buf.WriteString("/")
		buf.WriteString(k.KeyString())
		c.keysBuf[ki] = buf.Bytes()
	}
	for _, sop := range c.SpaceOps {
		for _, k := range sop.Add {
			addKey(sop.SpaceId, k, OpAdd)
		}
		for _, k := range sop.Delete {
			addKey(sop.SpaceId, k, OpDelete)
		}
		for _, k := range sop.Load {
			addKey(sop.SpaceId, k, OpLoad)
		}
	}
	return c.keysBuf
}

func (c *Cids) AddKey(key []byte) error {
	spaceId, op, k, err := parseCIDOp(key)
	if err != nil {
		return err
	}
	c.Add(spaceId, op, k)
	return nil
}

func (c *Cids) Len() (l int) {
	for _, so := range c.SpaceOps {
		l += len(so.Load)
		l += len(so.Delete)
		l += len(so.Add)
	}
	return
}

func (c *Cids) Release() {
	c.keysBuf = c.keysBuf[:0]
	for i, sop := range c.SpaceOps {
		c.SpaceOps[i].Add = sop.Add[:0]
		c.SpaceOps[i].Delete = sop.Delete[:0]
		c.SpaceOps[i].Load = sop.Load[:0]
	}
	c.SpaceOps = c.SpaceOps[:0]
	cidsPool.Put(c)
}