《游戏服务端编程实践》6.3.3 异步任务与保存队列

// file: ch06/save/queue.go
// tier: L1
// verify: go test ./ch06/save/

package save

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

var (
	ErrQueueClosed = errors.New("save queue closed")
	ErrQueueFull   = errors.New("save queue full")
)

type SaveTask struct {
	PlayerID uint64
	Data     PlayerData
	Priority bool
	Done     chan error
}

type SaveFunc func(ctx context.Context, task SaveTask) error

// SaveQueue 按 playerID 串行保存；同一玩家排队中的任务合并为最后一次快照。
type SaveQueue struct {
	save   SaveFunc
	maxLow int

	mu       sync.Mutex
	cond     *sync.Cond
	pending  map[uint64][]SaveTask
	active   map[uint64]bool
	inFlight int
	closed   bool
	wg       sync.WaitGroup
}

func NewSaveQueue(save SaveFunc, maxLowPending int) *SaveQueue {
	q := &SaveQueue{save: save, maxLow: maxLowPending, pending: map[uint64][]SaveTask{}, active: map[uint64]bool{}}
	q.cond = sync.NewCond(&q.mu)
	return q
}

func (q *SaveQueue) Enqueue(ctx context.Context, task SaveTask) error {
	if task.PlayerID == 0 {
		return errors.New("save task: player id required")
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	q.mu.Lock()
	defer q.mu.Unlock()
	if q.closed {
		return ErrQueueClosed
	}
	if !task.Priority && q.maxLow > 0 && q.outstandingLocked() >= q.maxLow {
		return ErrQueueFull
	}
	q.pending[task.PlayerID] = append(q.pending[task.PlayerID], task)
	if !q.active[task.PlayerID] {
		q.active[task.PlayerID] = true
		q.wg.Add(1)
		go q.worker(task.PlayerID)
	}
	return nil
}

func (q *SaveQueue) Flush(ctx context.Context) error {
	done := make(chan struct{})
	defer close(done)
	go func() {
		select {
		case <-ctx.Done():
			q.mu.Lock()
			q.cond.Broadcast()
			q.mu.Unlock()
		case <-done:
		}
	}()
	q.mu.Lock()
	defer q.mu.Unlock()
	for q.outstandingLocked() > 0 {
		if err := ctx.Err(); err != nil {
			return err
		}
		q.cond.Wait()
	}
	return nil
}

func (q *SaveQueue) Close(ctx context.Context) error {
	q.mu.Lock()
	if q.closed {
		q.mu.Unlock()
		return nil
	}
	q.closed = true
	q.mu.Unlock()
	if err := q.Flush(ctx); err != nil {
		return err
	}
	q.wg.Wait()
	return nil
}

func (q *SaveQueue) worker(playerID uint64) {
	defer q.wg.Done()
	for {
		q.mu.Lock()
		tasks := q.pending[playerID]
		if len(tasks) == 0 {
			q.active[playerID] = false
			delete(q.pending, playerID)
			q.cond.Broadcast()
			q.mu.Unlock()
			return
		}
		delete(q.pending, playerID)
		q.inFlight++
		merged := tasks[len(tasks)-1]
		q.mu.Unlock()

		err := q.save(context.Background(), merged)

		q.mu.Lock()
		q.inFlight--
		q.cond.Broadcast()
		q.mu.Unlock()
		deliver(tasks, err)
	}
}

func (q *SaveQueue) outstandingLocked() int {
	n := q.inFlight
	for _, tasks := range q.pending {
		n += len(tasks)
	}
	return n
}

func deliver(tasks []SaveTask, err error) {
	for _, task := range tasks {
		if task.Done == nil {
			continue
		}
		select {
		case task.Done <- err:
		default:
		}
	}
}

// file: ch06/save/queue_test.go
// tier: L1
// verify: go test ./ch06/save/

package save

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

func TestSaveQueueKeepsPerPlayerOrder(t *testing.T) {
	ctx := context.Background()
	var mu sync.Mutex
	var golds []int64
	calls := 0
	started := make(chan struct{})
	release := make(chan struct{})
	q := NewSaveQueue(func(_ context.Context, task SaveTask) error {
		mu.Lock()
		calls++
		call := calls
		golds = append(golds, task.Data.Gold)
		mu.Unlock()
		if call == 1 {
			close(started)
			<-release
		}
		return nil
	}, 0)
	first := make(chan error, 1)
	second := make(chan error, 1)
	if err := q.Enqueue(ctx, SaveTask{PlayerID: 1, Data: PlayerData{PlayerID: 1, Gold: 10}, Done: first}); err != nil {
		t.Fatal(err)
	}
	<-started
	if err := q.Enqueue(ctx, SaveTask{PlayerID: 1, Data: PlayerData{PlayerID: 1, Gold: 20}, Done: second}); err != nil {
		t.Fatal(err)
	}
	close(release)
	if err := <-first; err != nil {
		t.Fatal(err)
	}
	if err := <-second; err != nil {
		t.Fatal(err)
	}
	if err := q.Flush(ctx); err != nil {
		t.Fatal(err)
	}
	mu.Lock()
	defer mu.Unlock()
	if len(golds) != 2 || golds[0] != 10 || golds[1] != 20 {
		t.Fatalf("saves out of order: %v", golds)
	}
}

func TestSaveQueueMergesPending(t *testing.T) {
	ctx := context.Background()
	var mu sync.Mutex
	var golds []int64
	started := make(chan struct{})
	release := make(chan struct{})
	q := NewSaveQueue(func(_ context.Context, task SaveTask) error {
		mu.Lock()
		golds = append(golds, task.Data.Gold)
		first := len(golds) == 1
		mu.Unlock()
		if first {
			close(started)
			<-release
		}
		return nil
	}, 0)
	if err := q.Enqueue(ctx, SaveTask{PlayerID: 7, Data: PlayerData{PlayerID: 7, Gold: 1}}); err != nil {
		t.Fatal(err)
	}
	<-started
	for _, gold := range []int64{2, 3, 4} {
		if err := q.Enqueue(ctx, SaveTask{PlayerID: 7, Data: PlayerData{PlayerID: 7, Gold: gold}}); err != nil {
			t.Fatal(err)
		}
	}
	close(release)
	if err := q.Flush(ctx); err != nil {
		t.Fatal(err)
	}
	mu.Lock()
	defer mu.Unlock()
	if len(golds) != 2 || golds[0] != 1 || golds[1] != 4 {
		t.Fatalf("pending tasks not merged: %v", golds)
	}
}

func TestSaveQueueBackpressureAndClose(t *testing.T) {
	ctx := context.Background()
	started := make(chan struct{})
	release := make(chan struct{})
	q := NewSaveQueue(func(context.Context, SaveTask) error {
		close(started)
		<-release
		return nil
	}, 1)
	if err := q.Enqueue(ctx, SaveTask{PlayerID: 1, Data: PlayerData{PlayerID: 1}}); err != nil {
		t.Fatal(err)
	}
	<-started
	if err := q.Enqueue(ctx, SaveTask{PlayerID: 2, Data: PlayerData{PlayerID: 2}}); !errors.Is(err, ErrQueueFull) {
		t.Fatalf("want ErrQueueFull, got %v", err)
	}
	close(release)
	if err := q.Close(ctx); err != nil {
		t.Fatal(err)
	}
	if err := q.Enqueue(ctx, SaveTask{PlayerID: 3, Data: PlayerData{PlayerID: 3}}); !errors.Is(err, ErrQueueClosed) {
		t.Fatalf("want ErrQueueClosed, got %v", err)
	}
}

cd .local/checks/gs-practice-lab && gofmt -l ch06/save && go vet ./ch06/save/ && go test ./ch06/save/ -count=1

ok  	gspractice/ch06/save	0.115s