gyrogpsc/storage/kvstore.go

package storage

import (
	"encoding/binary"
	"encoding/json"
	"git.timovolkmann.de/gyrogpsc/core"
	"github.com/dgraph-io/badger/v2"
	"github.com/google/uuid"
	"github.com/sirupsen/logrus"
	"github.com/tidwall/pretty"
	"os"
	"path/filepath"
	"strconv"
)

// Must implement Repo
type badgerStore struct {
	trackingsDb *badger.DB
	recordsDb   *badger.DB
	rawdataDb   *badger.DB
}

func NewRepository(c *core.Configuration) *badgerStore {
	dir, _ := os.Getwd()
	logrus.Debug(dir)
	if _, err := os.Stat(filepath.Join(dir, "_db")); os.IsNotExist(err) {
		os.Mkdir(filepath.Join(dir, "_db"), os.ModePerm)
	}

	tr, err := badger.Open(badger.DefaultOptions("_db/trackings"))
	dp, err := badger.Open(badger.DefaultOptions("_db/records"))
	rd, err := badger.Open(badger.DefaultOptions("_db/raw"))

	if err != nil {
		logrus.Error(err)
	}
	return &badgerStore{trackingsDb: tr, recordsDb: dp, rawdataDb: rd}
}

func (r *badgerStore) isDbAvailable() bool {
	return r.trackingsDb.IsClosed() || r.recordsDb.IsClosed() || r.rawdataDb.IsClosed()
}

func (r *badgerStore) Save(tr core.Tracking) error {
	if ok := r.isDbAvailable(); ok {
		logrus.Error("unable to write to database. database closed!")
		return badger.ErrDBClosed
	}
	ts, err := tr.TimeCreated.MarshalText()
	if err != nil {
		logrus.Error(err, tr)
	}
	logrus.Info("save tracking:", tr.TimeCreated)
	meta, err := json.Marshal(tr.TrackingMetadata)
	if err != nil {
		logrus.Error(err, tr)
		return err
	}
	err = r.recordsDb.Update(func(txn *badger.Txn) error {
		for _, v := range tr.Records {
			k := createRecordKey(tr.UUID, v.RecordTimeKey.UnixNano())
			j, err := json.Marshal(v.DataPair)
			logrus.Debugln("save record k/v:\n", tr.UUID.String(), strconv.FormatInt(v.RecordTimeKey.UnixNano(), 10))
			logrus.Debugln(string(pretty.Pretty(j)))
			if err != nil {
				return err
			}
			txn.Set(k, j)
		}
		return nil
	})
	if err != nil {
		logrus.Error(err, tr)
		return err
	}
	err = r.rawdataDb.Update(func(txn *badger.Txn) error {
		for _, v := range tr.Rawdata {
			k := createRecordKey(tr.UUID, v.RecordTimeKey.UnixNano())
			j, err := json.Marshal(v)
			logrus.Debugln("save raw k/v:\n", tr.UUID.String(), strconv.FormatInt(v.RecordTimeKey.UnixNano(), 10))
			logrus.Debugln(string(pretty.Pretty(j)))
			if err != nil {
				return err
			}
			txn.Set(k, j)
		}
		return nil
	})
	if err != nil {
		logrus.Error(err, tr)
		return err
	}
	err = r.trackingsDb.Update(func(txn *badger.Txn) error {
		logrus.Debug("save tracking meta k/v:\n", string(ts), string(meta))
		err := txn.Set(ts, meta)
		return err
	})
	if err != nil {
		logrus.Error(err, tr)
		return err
	}

	dr := 0.5
	err = r.trackingsDb.RunValueLogGC(dr)
	logrus.Debug("DB GC:", err)
	err = r.recordsDb.RunValueLogGC(dr)
	logrus.Debug("DB GC:", err)
	err = r.rawdataDb.RunValueLogGC(dr)
	logrus.Debug("DB GC:", err)
	logrus.Info("sucessfully saved tracking")
	return nil
}

func (r *badgerStore) LoadAll() ([]core.TrackingMetadata, error) {
	var result []core.TrackingMetadata
	err := r.trackingsDb.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		it := txn.NewIterator(opts)
		defer it.Close()
		for it.Rewind(); it.Valid(); it.Next() {
			item := it.Item()
			el := core.TrackingMetadata{}
			item.Value(func(val []byte) error {
				logrus.Debugln(string(val))
				err := json.Unmarshal(val, &el)
				return err
			})
			result = append(result, el)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return result, nil
}

func (r *badgerStore) Load(id uuid.UUID) (core.Tracking, error) {
	panic("implement me")
}

func createRecordKey(uid uuid.UUID, timestamp int64) []byte {
	prefix, err := uid.MarshalText()
	if err != nil || timestamp < 0 {
		logrus.Error("unable to create key", err)
	}
	suffix := make([]byte, 8)
	binary.BigEndian.PutUint64(suffix, uint64(timestamp))

	return append(prefix, suffix...)
}

func unmarshalDataKey(key []byte) (uuid.UUID, int64) {
	if len(key) != 24 {
		panic("corrupted key")
	}
	prefix := key[0:15]
	suffix := key[15:24]
	uid, err := uuid.FromBytes(prefix)
	if err != nil {
		panic("corrupted key")
	}
	timestamp := int64(binary.BigEndian.Uint64(suffix))
	return uid, timestamp
}