jldb/mdb/testing/crashconsistency/main.go

package main

import (
	"crypto/rand"
	"errors"
	"hash/crc32"
	"log"
	mrand "math/rand"
	"os"
	"runtime"
	"slices"
	"sync"
	"sync/atomic"
	"time"

	"git.crumpington.com/public/jldb/mdb"
)

type DataItem struct {
	ID   uint64
	Data []byte
}

type DataCollection struct {
	*mdb.Collection[DataItem]
}

func NewDataCollection(db *mdb.Database) DataCollection {
	return DataCollection{
		Collection: mdb.NewCollection(db, "Data", &mdb.CollectionConfig[DataItem]{
			Copy: func(in *DataItem) *DataItem {
				out := new(DataItem)
				*out = *in
				out.Data = slices.Clone(in.Data)
				return out
			},
		}),
	}
}

type CRCItem struct {
	ID    uint64 // Always 1
	CRC32 uint32
}

type CRCCollection struct {
	*mdb.Collection[CRCItem]
}

func NewCRCCollection(db *mdb.Database) CRCCollection {
	return CRCCollection{
		Collection: mdb.NewCollection[CRCItem](db, "CRC", nil),
	}
}

type DataDB struct {
	*mdb.Database
	Datas DataCollection
	CRCs  CRCCollection
}

func OpenDataDB(rootDir string) (DataDB, error) {
	db := mdb.New(mdb.Config{RootDir: rootDir, Primary: true})
	testdb := DataDB{
		Database: db,
		Datas:    NewDataCollection(db),
		CRCs:     NewCRCCollection(db),
	}

	if err := db.Open(); err != nil {
		return testdb, err
	}

	return testdb, nil
}

func (db DataDB) ModifyFor(dt time.Duration) {
	wg := sync.WaitGroup{}
	var count int64
	for i := 0; i < runtime.NumCPU(); i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			t0 := time.Now()
			for time.Since(t0) < dt {
				atomic.AddInt64(&count, 1)
				db.modifyOnce()
			}
		}()
	}

	wg.Wait()
	log.Printf("Modified: %d", count)
}

func (db DataDB) modifyOnce() {
	isErr := mrand.Float64() < 0.1
	err := db.Update(func(tx *mdb.Snapshot) error {
		h := crc32.NewIEEE()
		for dataID := uint64(1); dataID < 10; dataID++ {
			d := DataItem{
				ID:   dataID,
				Data: make([]byte, 256),
			}

			rand.Read(d.Data)
			h.Write(d.Data)
			if err := db.Datas.Upsert(tx, &d); err != nil {
				return err
			}
		}

		crc := CRCItem{
			ID: 1,
		}

		if !isErr {
			crc.CRC32 = h.Sum32()
			return db.CRCs.Upsert(tx, &crc)
		}

		crc.CRC32 = 1
		if err := db.CRCs.Upsert(tx, &crc); err != nil {
			return err
		}

		return errors.New("ERROR")
	})

	if isErr != (err != nil) {
		panic(err)
	}
}

func (db DataDB) ComputeCRC(tx *mdb.Snapshot) uint32 {
	h := crc32.NewIEEE()
	for dataID := uint64(1); dataID < 10; dataID++ {
		d, ok := db.Datas.ByID.Get(tx, &DataItem{ID: dataID})
		if !ok {
			continue
		}
		h.Write(d.Data)
	}
	return h.Sum32()
}

func (db DataDB) ReadCRC(tx *mdb.Snapshot) uint32 {
	r, ok := db.CRCs.ByID.Get(tx, &CRCItem{ID: 1})
	if !ok {
		return 0
	}
	return r.CRC32
}

func main() {
	db, err := OpenDataDB(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}

	db.ModifyFor(time.Minute)
}