《MMO 服务器开发》9.2 落库与脏写防护

// file: internal/store/persist.go
// tier: L2
// verify: go test ./internal/store/ -run TestOptimisticLockDirtyWrite -v

package store

import (
	"context"
	"errors"
	"fmt"

	"github.com/jackc/pgx/v5/pgxpool"
)

// ErrStaleWrite is returned when an UPDATE matches zero rows because the row's
// version moved on since this writer read it. The caller must reload and
// reapply its change; blindly writing would clobber the other writer.
var ErrStaleWrite = errors.New("store: stale write (version mismatch)")

// DirtyMask records which tables a flush must touch. Extension tables are only
// rewritten when their bit is set, so a position change never rewrites the bag
// and a quest turn-in never rewrites the attribute bag.
type DirtyMask uint8

const (
	DirtyMain DirtyMask = 1 << iota
	DirtyItems
	DirtyQuests
	DirtyAttrs
)

// Flush persists a player in one transaction. The main UPDATE carries the
// concurrency gate: `WHERE player_id = $1 AND version = $8`. If two writers
// race, the loser's WHERE matches nothing, it gets ErrStaleWrite, and it never
// touches any extension table. The version is bumped in the same statement, so
// the next writer must have read the row after this commit to win.
func Flush(ctx context.Context, pool *pgxpool.Pool, p *Player,
	items []Item, quests []Quest, attrs []Attr, dirty DirtyMask) error {

	tx, err := pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)

	if dirty&DirtyMain != 0 {
		ct, err := tx.Exec(ctx, fmt.Sprintf(`UPDATE %s.player SET
			account=$2, name=$3, level=$4, pos_x=$5, pos_y=$6, gold=$7,
			version = version + 1, updated_at = now()
			WHERE player_id=$1 AND version=$8`, Schema),
			p.ID, p.Account, p.Name, p.Level, p.PosX, p.PosY, p.Gold, p.Version)
		if err != nil {
			return err
		}
		if ct.RowsAffected() == 0 {
			return ErrStaleWrite
		}
	}

	if dirty&DirtyItems != 0 {
		if _, err := tx.Exec(ctx, fmt.Sprintf(
			"DELETE FROM %s.player_item WHERE player_id=$1", Schema), p.ID); err != nil {
			return err
		}
		for _, it := range items {
			if _, err := tx.Exec(ctx, fmt.Sprintf(`INSERT INTO %s.player_item
				(inst_id, player_id, item_id, count) VALUES ($1,$2,$3,$4)`, Schema),
				it.InstID, it.PlayerID, it.ItemID, it.Count); err != nil {
				return err
			}
		}
	}

	if dirty&DirtyQuests != 0 {
		for _, q := range quests {
			if _, err := tx.Exec(ctx, fmt.Sprintf(`INSERT INTO %s.player_quest
				(player_id, quest_id, state) VALUES ($1,$2,$3)
				ON CONFLICT (player_id, quest_id) DO UPDATE SET state=EXCLUDED.state`, Schema),
				q.PlayerID, q.QuestID, q.State); err != nil {
				return err
			}
		}
	}

	if dirty&DirtyAttrs != 0 {
		for _, a := range attrs {
			if _, err := tx.Exec(ctx, fmt.Sprintf(`INSERT INTO %s.player_attr
				(player_id, key, value) VALUES ($1,$2,$3)
				ON CONFLICT (player_id, key) DO UPDATE SET value=EXCLUDED.value`, Schema),
				a.PlayerID, a.Key, a.Value); err != nil {
				return err
			}
		}
	}

	if err := tx.Commit(ctx); err != nil {
		return err
	}
	// The DB bumped the version; keep the in-memory copy in step so a second
	// flush from the same object is not immediately stale.
	p.Version++
	return nil
}

// DirtyTracker accumulates which tables changed since the last flush. It is
// the write-combining half of the story: a hundred ticks that only move the
// player collapse into one DirtyMain flush.
type DirtyTracker struct {
	mask DirtyMask
}

func (d *DirtyTracker) Mark(m DirtyMask) { d.mask |= m }
func (d *DirtyTracker) Mask() DirtyMask  { return d.mask }
func (d *DirtyTracker) Reset()           { d.mask = 0 }
func (d *DirtyTracker) Dirty() bool      { return d.mask != 0 }

// file: internal/store/persist_test.go
// tier: L2
// verify: go test ./internal/store/ -run TestOptimisticLockDirtyWrite -v

package store

import (
	"context"
	"errors"
	"sync"
	"testing"
)

// TestOptimisticLockDirtyWrite proves the two properties of the version gate:
// when two writers race on the same row, exactly one wins, and the version is
// bumped by exactly one so the loser can detect the collision and retry.
func TestOptimisticLockDirtyWrite(t *testing.T) {
	pool := testPool(t)
	defer pool.Close()
	resetSchema(t, pool)
	ctx := context.Background()

	if err := SavePlayer(ctx, pool, &Player{
		ID: 1001, Account: "acc-1001", Name: "晨光", Level: 10, Gold: 500, Version: 1,
	}); err != nil {
		t.Fatal(err)
	}

	// Both writers read the same row, so both hold version 1.
	a, err := LoadPlayer(ctx, pool, 1001)
	if err != nil {
		t.Fatal(err)
	}
	b, err := LoadPlayer(ctx, pool, 1001)
	if err != nil {
		t.Fatal(err)
	}
	if a.Version != 1 || b.Version != 1 {
		t.Fatalf("both should read version 1, got %d/%d", a.Version, b.Version)
	}
	a.Gold = 600
	b.Gold = 700

	var wg sync.WaitGroup
	errs := make([]error, 2)
	wg.Add(2)
	go func() { defer wg.Done(); errs[0] = Flush(ctx, pool, a, nil, nil, nil, DirtyMain) }()
	go func() { defer wg.Done(); errs[1] = Flush(ctx, pool, b, nil, nil, nil, DirtyMain) }()
	wg.Wait()

	success, stale := 0, 0
	for _, e := range errs {
		switch {
		case e == nil:
			success++
		case errors.Is(e, ErrStaleWrite):
			stale++
		default:
			t.Fatalf("unexpected error: %v", e)
		}
	}
	if success != 1 || stale != 1 {
		t.Fatalf("success=%d stale=%d, want exactly 1 and 1", success, stale)
	}

	final, err := LoadPlayer(ctx, pool, 1001)
	if err != nil {
		t.Fatal(err)
	}
	if final.Version != 2 {
		t.Fatalf("version=%d want 2 (bumped exactly once)", final.Version)
	}
	// The surviving value is one writer's whole value, never a merge.
	if final.Gold != 600 && final.Gold != 700 {
		t.Fatalf("gold=%d, want 600 or 700", final.Gold)
	}
	t.Logf("version gate ok: 1 commit, 1 stale, version=2, gold=%d", final.Gold)
}