gpu-compute-chain/x/cosmostest/memdb/biddb.go

package memdb

import (
	"bytes"
	"cosmos-test/x/cosmostest/types"
	"encoding/gob"
	"errors"
	"log"
	"math/big"

	"github.com/dgraph-io/badger/v3"
)

type bidDB struct {
	db *badger.DB
}

var BidDB bidDB

// Mount Db & initialize encoder/decoder
func (b *bidDB) Mount() {
	db, err := badger.Open(badger.DefaultOptions("").WithInMemory(true))
	if err != nil {
		// must force crash here, since db is absolutely required
		log.Fatalf("Failed to mount in-memory db: %s", err)
	}
	b.db = db

	// initialize encode/decode stuff
	gob.Register([]types.Bid{})
}

// Add a bid to the bid list under specified auction key.
func (b *bidDB) AddBid(auctionId string, bid *types.Bid) error {
	k := []byte(auctionId)

	err := b.db.Update(func(txn *badger.Txn) error {
		var bids []*types.Bid
		bidsOld, err := txn.Get(k)
		if errors.Is(err, badger.ErrKeyNotFound) {
			// key not found -> just create a new Bid array
			bids = []*types.Bid{}
		} else {
			if err != nil {
				return err
			}
			// key found -> decode contents to bids
			bidsOld.Value(func(val []byte) error {
				dec := gob.NewDecoder(bytes.NewReader(val))
				if err := dec.Decode(&bids); err != nil {
					return err
				}
				return nil
			})
		}

		// append bid
		bids = append(bids, bid)
		// encode new list
		buf := bytes.NewBuffer(nil)
		enc := gob.NewEncoder(buf)
		if err := enc.Encode(&bids); err != nil {
			return err
		}
		// set under auction key in db
		if err := txn.Set(k, buf.Bytes()); err != nil {
			return err
		}

		return nil
	})

	return err
}

// Get the highest bid in the list under specified auction key.
func (b *bidDB) GetLowestBid(auctionId string) (*types.Bid, error) {
	k := []byte(auctionId)

	var bid *types.Bid

	err := b.db.View(func(txn *badger.Txn) error {
		bidData, err := txn.Get(k)
		if err != nil {
			if !errors.Is(err, badger.ErrKeyNotFound) {
				return err
			} else {
				return nil
			}
		}

		err = bidData.Value(func(val []byte) error {
			var bids []*types.Bid
			dec := gob.NewDecoder(bytes.NewReader(val))
			if err := dec.Decode(&bids); err != nil {
				return err
			}

			if bids == nil {
				return nil
			}

			minAmt := big.NewInt(0)
			refAmt := big.NewInt(0)
			for _, currBid := range bids {
				if bid == nil {
					bid = currBid
					minAmt.SetString(currBid.Amount, 10)
				} else {
					// set ref amt.
					refAmt.SetString(currBid.Amount, 10)
					// if ref amt is less, then we have a new min
					if refAmt.Cmp(minAmt) == -1 {
						bid = currBid
					}
				}
			}
			return nil
		})

		return err
	})

	return bid, err
}

func (b *bidDB) GetBids(auctionId string) ([]*types.Bid, error) {
	k := []byte(auctionId)

	var bids []*types.Bid
	err := b.db.View(func(txn *badger.Txn) error {
		res, err := txn.Get(k)
		if err != nil {
			return err
		} else {
			err := res.Value(func(val []byte) error {
				dec := gob.NewDecoder(bytes.NewReader(val))
				err := dec.Decode(&bids)
				return err
			})
			return err
		}
	})

	return bids, err
}

func (b *bidDB) ClearAuction(auctionId string) error {
	k := []byte(auctionId)

	err := b.db.Update(func(txn *badger.Txn) error {
		return txn.Delete(k)
	})

	return err
}

// Iterate over all auction keys in memory. VIEW-ONLY.
func (b *bidDB) ForEachAuction(viewFunc func(string) error) error {
	err := b.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		// can customize options down here if we want
		iter := txn.NewIterator(opts)
		defer iter.Close()
		for iter.Rewind(); iter.Valid(); iter.Next() {
			item := iter.Item()
			key := string(item.Key())
			err := viewFunc(key)
			if err != nil {
				return err
			}
		}
		return nil
	})
	return err
}