Vlog GC example

main_test.go

package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"path/filepath"
	"syscall"
	"testing"
	"time"

	"github.com/dustin/go-humanize"
	ds "github.com/ipfs/go-datastore"

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

const bdir = "/tmp/badger"

const (
	entryC = 300000
	preC   = 10000

	entryS = 10000
)

func TestGc(t *testing.T) {
	opts := badger.DefaultOptions

	if err := os.RemoveAll(bdir); err != nil {
		t.Fatal(err)
	}

	opts.Dir = bdir
	opts.ValueDir = bdir
	opts.Logger = nil

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

	r := rand.New(rand.NewSource(555))

	wb := db.NewWriteBatch()

	for i := 0; i < preC; i++ { // put non-deletable entries
		b, err := ioutil.ReadAll(io.LimitReader(r, entryS))
		if err != nil {
			t.Fatal(err)
		}
		if err := wb.Set(ds.RandomKey().Bytes(), b, 0); err != nil {
			t.Fatal(err)
		}
	}

	if err := wb.Flush(); err != nil {
		t.Fatal(err)
	}

	pds(t, "non-deletable put")

	wb = db.NewWriteBatch()
	es := make([][]byte, entryC)
	for i := 0; i < entryC; i++ { // put deletable entries
		b, err := ioutil.ReadAll(io.LimitReader(r, entryS))
		if err != nil {
			t.Fatal(err)
		}
		es[i] = ds.RandomKey().Bytes()
		if err := wb.Set(es[i], b, 0); err != nil {
			t.Fatal(err)
		}
	}

	if err := wb.Flush(); err != nil {
		t.Fatal(err)
	}

	pds(t, "data put")

	if err := db.Close(); err != nil {
		t.Fatal(err)
	}

	pds(t, "put closed")

	db, err = badger.Open(opts)
	if err != nil {
		t.Fatal(err)
	}

	pds(t, "del-open")

	wb = db.NewWriteBatch()
	for _, e := range es {
		if err := wb.Delete(e); err != nil {
			t.Fatal(err)
		}
	}
	if err := wb.Flush(); err != nil {
		t.Fatal(err)
	}

	/*
	* The following code was added so that GC can remove the entries added via batch API
	 */
	txn := db.NewTransaction(true)
	txn.Set([]byte("key"), []byte("value"))
	txn.Commit()
	// THIS IS IMPORTANT
	// We calculate the amount of data that can be removed when compactions are triggered.
	// In this case, we force compaction by closing the database
	if err := db.Close(); err != nil {
		t.Fatal(err)
	}

	pds(t, "after del")

	for err == nil {
		err = db.RunValueLogGC(0.5)
		pds(t, "gc-step")
	}
	if err != badger.ErrNoRewrite {
		t.Fatal(err)
	}

	pds(t, "gc")

	if err := db.Close(); err != nil {
		t.Fatal(err)
	}
	pds(t, "close-gc")
}

func pds(t *testing.T, prefix string) {
	time.Sleep(1 * time.Second) // wait for things to settle down a bit

	i, real, sst, vlog, err := dirSize(bdir)
	if err != nil {
		t.Fatal(err)
	}

	fmt.Printf("\x1b[32m%s: %s(%s); sst: %s; vlog: %s\x1b[39m\n",
		prefix,
		humanize.Bytes(uint64(i)),
		humanize.Bytes(uint64(real)),
		humanize.Bytes(uint64(sst)),
		humanize.Bytes(uint64(vlog)))
}

func dirSize(path string) (int64, int64, int64, int64, error) {
	var size int64
	var real int64
	var sst int64
	var vlog int64

	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() {
			size += info.Size()
			real += (info.Sys().(*syscall.Stat_t)).Blocks * 512

			if len(info.Name()) < 10 {
				return err
			}
			if info.Name()[7] == 's' {
				sst += info.Size()
			}
			if info.Name()[7] == 'v' {
				vlog += info.Size()
			}
		}
		return err
	})
	return size, real, sst, vlog, err
}

output.txt

=== RUN   TestGc
non-deletable put: 101 MB(101 MB); sst: 0 B; vlog: 101 MB
data put: 3.1 GB(3.1 GB); sst: 14 MB; vlog: 3.1 GB
put closed: 3.1 GB(3.1 GB); sst: 20 MB; vlog: 3.1 GB
del-open: 3.1 GB(3.1 GB); sst: 20 MB; vlog: 3.1 GB
after del: 3.2 GB(3.2 GB); sst: 17 MB; vlog: 3.1 GB
gc-step: 1.2 GB(1.2 GB); sst: 17 MB; vlog: 1.2 GB
gc-step: 136 MB(136 MB); sst: 17 MB; vlog: 120 MB
gc-step: 136 MB(136 MB); sst: 17 MB; vlog: 120 MB
gc: 136 MB(136 MB); sst: 17 MB; vlog: 120 MB
close-gc: 137 MB(137 MB); sst: 17 MB; vlog: 120 MB
--- PASS: TestGc (29.72s)
PASS
ok  	github.com/jarifibrahim/foo	30.085s
Success: Tests passed.