gocrew3 sqlite

main.go

package main

import (
	"database/sql"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"time"

	_ "github.com/glebarez/go-sqlite"
)

func main() {
	if err := run(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func run() error {
	data := generateRandomData(1_000_000)
	dir, err := os.MkdirTemp("", "test-")
	if err != nil {
		return err
	}

	defer os.RemoveAll(dir)

	fn := filepath.Join(dir, "db")

	db, err := sql.Open("sqlite", fn)
	if err != nil {
		return err
	}

	_, err = db.Exec(`PRAGMA page_size=32768`)
	if err != nil {
		panic(err)
	}

	_, err = db.Exec(`PRAGMA synchronous = NORMAL`)
	if err != nil {
		panic(err)
	}

	_, err = db.Exec(`PRAGMA journal_mode = WAL`)
	if err != nil {
		panic(err)
	}

	if _, err = db.Exec(createTableQuery); err != nil {
		return err
	}

	cnt := insert(db, data)
	runSelect(db, data[:cnt])
	runUpdate(db, data[:cnt])

	if err = db.Close(); err != nil {
		return err
	}

	return nil
}

func insert(db *sql.DB, data []TestItem) int {
	signalCh := make(chan struct{})
	retCh := make(chan int)

	go func() {
		cnt := 0
		i := 0
		stmt, err := db.Prepare("insert into test values(?,?,?,?,?,?,?,?,?)")
		if err != nil {
			panic(err)
		}

		for {
			select {
			case <-signalCh:
				retCh <- cnt
				return
			default:
				d := data[i]
				_, err := stmt.Exec(d.FieldID, d.FileName, d.Size, d.Locked, d.BucketID, d.TTL, d.SHA, d.Secret, d.Namespace)
				if err != nil {
					panic(err)
				}
				cnt++
				i++
				if i == len(data) {
					i = 0
				}
			}
		}
	}()

	time.Sleep(time.Second)
	signalCh <- struct{}{}
	cnt := <-retCh
	fmt.Println("Insert", cnt)

	row := db.QueryRow("select count(*) from test")
	if row.Err() != nil {
		panic(row.Err())
	}

	var count int
	err := row.Scan(&count)
	if err != nil {
		panic(err)
	}
	if count != cnt {
		panic("diff count")
	}

	return cnt
}

func runSelect(db *sql.DB, data []TestItem) {
	signalCh := make(chan struct{})
	retCh := make(chan int)

	go func() {
		var targetSHA string
		stmt, err := db.Prepare("select sha from test where file_id=?")
		if err != nil {
			panic(err)
		}

		rand.Shuffle(len(data), func(i, j int) {
			data[i], data[j] = data[j], data[i]
		})

		cnt := 0
		i := 0
		for {
			select {
			case <-signalCh:
				retCh <- cnt
				return
			default:
				d := data[i]
				err := stmt.QueryRow(d.FieldID).Scan(&targetSHA)
				if err != nil {
					panic(err)
				}

				if d.SHA != targetSHA {
					panic("diff sha")
				}

				cnt++
				i++
				if i >= len(data) {
					i = 0
				}
			}
		}
	}()
	time.Sleep(time.Second)
	signalCh <- struct{}{}
	cnt := <-retCh

	fmt.Println("Select", cnt)
}

func runUpdate(db *sql.DB, data []TestItem) {
	signalCh := make(chan struct{})
	retCh := make(chan int)

	go func() {
		stmt, err := db.Prepare("update test SET sha=? where file_id=?")
		if err != nil {
			panic(err)
		}

		rand.Shuffle(len(data), func(i, j int) {
			data[i], data[j] = data[j], data[i]
		})

		cnt := 0
		i := 0
		for {
			select {
			case <-signalCh:
				retCh <- cnt
				return
			default:
				d := data[i]
				_, err := stmt.Exec(randString(28), d.FieldID)
				if err != nil {
					panic(err)
				}
				cnt++
				i++
				if i == len(data) {
					i = 0
				}
			}
		}
	}()
	time.Sleep(time.Second)
	signalCh <- struct{}{}
	cnt := <-retCh

	fmt.Println("Update", cnt)
}

type TestItem struct {
	FieldID                int
	FileName               string
	Size                   int
	Locked                 int
	BucketID               int
	TTL                    int
	SHA, Secret, Namespace string
}

func generateRandomData(size int) []TestItem {
	start := time.Now()
	ret := make([]TestItem, size)
	for i := 0; i < size; i++ {
		ret[i] = TestItem{
			FieldID:   rand.Int(),
			FileName:  randString(rand.Intn(900) + 10),
			Size:      rand.Int(),
			Locked:    boolToInt(rand.Float64() > 0.5),
			BucketID:  rand.Int(),
			TTL:       rand.Int(),
			SHA:       randString(28),
			Secret:    randString(16),
			Namespace: randString(10),
		}
	}
	fmt.Println("Data generated ", time.Since(start))
	return ret
}

func boolToInt(b bool) int {
	if b {
		return 1
	}
	return 0
}

const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

func randString(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = letterBytes[rand.Intn(len(letterBytes))]
	}
	return string(b)
}

const createTableQuery = `
drop table if exists test;
create table test(file_id int, file_name text, size int, locked int, bucket_id int, ttl int, sha text, secret text, namespace text);
create index test_file_id_idx on test(file_id);
`