NixMsg/internal/store/queue_test.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"sync"
	"testing"
	"time"
)

func TestWriteQueueSavepointIsolatesFailure(t *testing.T) {
	t.Parallel()
	dir := t.TempDir()
	db, err := Open(dir, "FULL")
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = db.Close() }()

	ctx := context.Background()
	var wg sync.WaitGroup
	wg.Add(2)
	errOK := make(chan error, 1)
	errBad := make(chan error, 1)

	go func() {
		defer wg.Done()
		errOK <- db.Queue.Do(ctx, func(tx *sql.Tx) error {
			_, e := tx.Exec(
				`INSERT INTO settings(key, value, updated_at) VALUES(?, ?, ?)`,
				"k_ok", "1", time.Now().UnixMilli(),
			)
			return e
		})
	}()
	go func() {
		defer wg.Done()
		// 稍等，尽量与成功操作进同一批。
		time.Sleep(500 * time.Microsecond)
		errBad <- db.Queue.Do(ctx, func(tx *sql.Tx) error {
			return errors.New("forced op failure")
		})
	}()
	wg.Wait()

	if e := <-errOK; e != nil {
		t.Fatalf("ok op: %v", e)
	}
	if e := <-errBad; e == nil || e.Error() != "forced op failure" {
		t.Fatalf("bad op: %v", e)
	}

	var value string
	if scanErr := db.Read.QueryRow(`SELECT value FROM settings WHERE key = ?`, "k_ok").Scan(&value); scanErr != nil {
		t.Fatalf("ok row missing: %v", scanErr)
	}
	if value != "1" {
		t.Fatalf("value=%q", value)
	}
	if !db.Queue.IsReady() {
		t.Fatal("op failure should not mark queue busy")
	}
}

func TestWriteQueueBatchCommit(t *testing.T) {
	t.Parallel()
	dir := t.TempDir()
	db, err := Open(dir, "FULL")
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = db.Close() }()

	ctx := context.Background()
	const n = 32
	var wg sync.WaitGroup
	wg.Add(n)
	errs := make([]error, n)
	for i := 0; i < n; i++ {
		i := i
		go func() {
			defer wg.Done()
			key := fmt.Sprintf("batch_%d", i)
			errs[i] = db.Queue.Do(ctx, func(tx *sql.Tx) error {
				_, e := tx.Exec(
					`INSERT INTO settings(key, value, updated_at) VALUES(?, ?, ?)`,
					key, "1", time.Now().UnixMilli(),
				)
				return e
			})
		}()
	}
	wg.Wait()
	for i, e := range errs {
		if e != nil {
			t.Fatalf("op %d: %v", i, e)
		}
	}
	var count int
	if err := db.Read.QueryRow(`SELECT COUNT(*) FROM settings WHERE key LIKE 'batch_%'`).Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != n {
		t.Fatalf("count=%d want %d", count, n)
	}
}

func TestWriteQueueNoBacklogAt200PerSec(t *testing.T) {
	t.Parallel()
	dir := t.TempDir()
	db, err := Open(dir, "FULL")
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = db.Close() }()

	ctx := context.Background()
	const total = 200
	start := time.Now()
	var wg sync.WaitGroup
	wg.Add(total)
	for i := 0; i < total; i++ {
		i := i
		go func() {
			defer wg.Done()
			key := fmt.Sprintf("load_%d", i)
			if e := db.Queue.Do(ctx, func(tx *sql.Tx) error {
				_, e := tx.Exec(
					`INSERT INTO settings(key, value, updated_at) VALUES(?, ?, ?)`,
					key, "1", time.Now().UnixMilli(),
				)
				return e
			}); e != nil {
				t.Errorf("op %d: %v", i, e)
			}
		}()
	}
	wg.Wait()
	elapsed := time.Since(start)
	if db.Queue.Len() != 0 {
		t.Fatalf("queue backlog len=%d", db.Queue.Len())
	}
	// 200 条并发写入应在数秒内完成（合并提交）；过长则合并未生效。
	if elapsed > 5*time.Second {
		t.Fatalf("200 writes took %s, queue likely not merging well", elapsed)
	}
}

func BenchmarkWriteQueue200PerSec(b *testing.B) {
	dir := b.TempDir()
	db, err := Open(dir, "FULL")
	if err != nil {
		b.Fatal(err)
	}
	defer func() { _ = db.Close() }()

	ctx := context.Background()
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		key := fmt.Sprintf("bench_%d", i)
		if err := db.Queue.Do(ctx, func(tx *sql.Tx) error {
			_, e := tx.Exec(
				`INSERT INTO settings(key, value, updated_at) VALUES(?, ?, ?)
ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated_at = excluded.updated_at`,
				key, "1", time.Now().UnixMilli(),
			)
			return e
		}); err != nil {
			b.Fatal(err)
		}
	}
}