测试sqlite作为网站数据库性能

sqlite_site_test.go

package main

import (
	"context"
	cRand "crypto/rand"
	"database/sql"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	mRand "math/rand"
	"net/http"
	"os"
	"time"

	// driver for sqlite3
	_ "github.com/mattn/go-sqlite3"
)

var (
	db         *sql.DB
	dbFile     string
	listenAddr string
	loopNumber int
)

func init() {
	flag.StringVar(&dbFile, "db", "", "Path to the test database file")
	flag.StringVar(&listenAddr, "listen", "127.0.0.1:12345", "Address to listen on")
	flag.IntVar(&loopNumber, "loop", 2, "Number of loops")
	flag.Parse()

	if dbFile == "" || loopNumber <= 0 {
		flag.Usage()
		os.Exit(1)
	}

	dsn := fmt.Sprintf("file:%s?mode=rwc&_timeout=5000&_journal=WAL&_sync=NORMAL", dbFile)

	var err error
	db, err = sql.Open("sqlite3", dsn)
	if err != nil {
		panic(fmt.Errorf("open database, %w", err))
	}

	_, err = db.Exec(`create table if not exists test (
		id integer primary key,
		rand_number integer not null,
		rand_blob blob not null,
		created_at int not null
	)`)
	if err != nil {
		panic(fmt.Errorf("create table, %w", err))
	}

	for range 100 {
		if row, err := newRow(); err != nil {
			panic(fmt.Errorf("generate new row, %w", err))
		} else if err := insertRow(row, db); err != nil {
			panic(fmt.Errorf("insert new row, %w", err))
		}
	}
}

func main() {
	defer db.Close()

	mux := http.NewServeMux()

	mux.HandleFunc("GET /write", func(w http.ResponseWriter, r *http.Request) {
		for range loopNumber {
			row, err := newRow()
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}

			if err := insertRow(row, db); err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}
		}

		w.WriteHeader(http.StatusCreated)
	})

	mux.HandleFunc("GET /transaction", func(w http.ResponseWriter, r *http.Request) {
		if err := runTransaction(db, func(tx *sql.Tx) error {
			for range loopNumber {
				row, err := newRow()
				if err != nil {
					return fmt.Errorf("generate new row, %w", err)
				}

				if err := insertRow(row, tx); err != nil {
					return fmt.Errorf("insert new row, %w", err)
				}
			}

			return nil
		}); err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}

		w.WriteHeader(http.StatusCreated)
	})

	mux.HandleFunc("GET /read", func(w http.ResponseWriter, r *http.Request) {
		for range loopNumber {
			id := mRand.Intn(100) + 1

			row, err := selectRow(r.Context(), id)
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}

			_ = json.NewEncoder(w).Encode(row)
		}
	})

	fmt.Println("listening on", listenAddr)
	if err := http.ListenAndServe(listenAddr, mux); err != nil && err != http.ErrServerClosed {
		panic(fmt.Errorf("start http server, %w", err))
	}
}

func newRow() (map[string]any, error) {
	rs := make([]byte, 512)
	if _, err := cRand.Read(rs); err != nil {
		return nil, fmt.Errorf("read random bytes, %w", err)
	}

	row := map[string]any{
		"rand_number": mRand.Intn(10000),
		"rand_blob":   rs,
		"created_at":  time.Now().Unix(),
	}
	return row, nil
}

type execer interface {
	Exec(query string, args ...any) (sql.Result, error)
}

func insertRow(row map[string]any, db execer) error {
	_, err := db.Exec(`insert into test (rand_number, rand_blob, created_at) values (?, ?, ?)`,
		row["rand_number"], row["rand_blob"], row["created_at"])
	return err
}

func selectRow(ctx context.Context, id int) (map[string]any, error) {
	res := db.QueryRowContext(ctx, `select rand_number, rand_blob, created_at from test where id = ?`, id)

	var (
		randNumber int
		randBlob   []byte
		createdAt  int64
	)

	if err := res.Scan(&randNumber, &randBlob, &createdAt); err != nil {
		return nil, fmt.Errorf("scan row, %w", err)
	}
	return map[string]any{
		"id":          id,
		"rand_number": randNumber,
		"rand_blob":   randBlob,
		"created_at":  createdAt,
	}, nil
}

func runTransaction(db *sql.DB, fn func(tx *sql.Tx) error) (err error) {
	tx, err := db.Begin()
	if err != nil {
		return err
	}

	defer func() {
		if err == nil {
			err = tx.Commit()
		} else {
			err = errors.Join(err, tx.Rollback())
		}
	}()

	return fn(tx)
}