badger_write_skew.go

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strconv"

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

func mustParseNum(v []byte) int {
	n, err := strconv.ParseInt(string(v), 10, 64)
	if err != nil {
		log.Fatal(err)
	}
	return int(n)
}

func getCounts(it *badger.Iterator) (aggregateEvens, aggregateOdds, evens, odds int) {
	for it.Rewind(); it.Valid(); it.Next() {
		if it.ValidForPrefix([]byte{'_'}) {
			v, err := it.Item().Value()
			if err != nil {
				log.Fatal(err)
			}
			n := mustParseNum(v)
			if string(it.Item().Key()) == "_even" {
				aggregateEvens = n
			} else {
				aggregateOdds = n
			}
			continue
		}
		n := mustParseNum(it.Item().Key())
		if n%2 == 0 {
			evens++
		} else {
			odds++
		}
	}
	return
}

func main() {
	dir, err := ioutil.TempDir("", "badger")
	if err != nil {
		log.Fatal(err)
	}
	defer os.RemoveAll(dir)

	opts := badger.DefaultOptions
	opts.Dir = dir
	opts.ValueDir = dir

	db, err := badger.Open(opts)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	bkey := func(i int) []byte {
		return []byte(fmt.Sprintf("%d", i))
	}

	// Initialize
	txn := db.NewTransaction(true)

	// Set keys 0, 2, 4
	for i := 0; i <= 4; i += 2 {
		err := txn.Set(bkey(i), []byte{})
		if err != nil {
			log.Fatal(err)
		}
	}
	err = txn.Commit(nil)
	if err != nil {
		log.Fatal(err)
	}

	// Now start two concurrent transactions.

	txnA := db.NewTransaction(true)
	txnB := db.NewTransaction(true)

	{
		// txnA adds an even number to the set and updates _odd aggregate count
		err = txnA.Set(bkey(6), []byte{})
		if err != nil {
			log.Fatal(err)
		}
		it := txnA.NewIterator(badger.DefaultIteratorOptions)
		_, _, _, odds := getCounts(it)
		// Update count
		err = txnA.Set([]byte("_odd"), []byte(fmt.Sprintf("%d", odds)))
		if err != nil {
			log.Fatal(err)
		}
	}

	{
		// txnB adds an odd number to the set and updates _even aggregate count
		err = txnB.Set(bkey(1), []byte{})
		if err != nil {
			log.Fatal(err)
		}
		it := txnB.NewIterator(badger.DefaultIteratorOptions)
		_, _, evens, _ := getCounts(it)
		// Update count
		err = txnB.Set([]byte("_even"), []byte(fmt.Sprintf("%d", evens)))
		if err != nil {
			log.Fatal(err)
		}
	}

	err = txnA.Commit(nil)
	if err != nil {
		log.Fatal(err)
	}

	err = txnB.Commit(nil)
	if err != nil {
		log.Fatal(err)
	}

	err = db.View(func(txn *badger.Txn) error {
		it := txn.NewIterator(badger.DefaultIteratorOptions)
		aggregateEvens, aggregateOdds, evens, odds := getCounts(it)
		if aggregateEvens != evens || aggregateOdds != odds {
			return fmt.Errorf("Total (evens, odds) = (%d, %d) doesn't match (_even, _odd) = (%d, %d)",
				evens, odds, aggregateEvens, aggregateOdds)
		}

		return nil
	})
	if err != nil {
		log.Fatal(err)
	}
}