perf(cmpp): instrument inbound flow and unbatch submit worker

2026-08-20 11:27:35 +08:00
parent c4f36fc50d
commit b9a71fe0b9
23 changed files with 760 additions and 138 deletions
@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"

 	"cmpp-platform/gateway/internal/metrics"
@@ -21,11 +22,12 @@ import (
 )

 const (
-	defaultStream   = "gateway.submit.commands"
-	defaultGroup    = "cmpp-gateway"
-	defaultConsumer = "gateway-1"
-	defaultMinIdle  = 30 * time.Second
-	defaultMaxFails = 3
+	defaultStream      = "gateway.submit.commands"
+	defaultGroup       = "cmpp-gateway"
+	defaultConsumer    = "gateway-1"
+	defaultMinIdle     = 30 * time.Second
+	defaultMaxFails    = 3
+	defaultConcurrency = 64
 )

 type Worker struct {
@@ -39,11 +41,13 @@ type Worker struct {
 	Consumer         string
 	Block            time.Duration
 	Count            int64
+	Concurrency      int
 	MinIdle          time.Duration
 	MaxFailures      int
 	APIBaseURL       string
 	HTTPClient       *http.Client
 	Logger           *log.Logger
+	inFlight         atomic.Int64
 }

 type DeadLetterEvent struct {
@@ -78,6 +82,8 @@ func (w *Worker) Run(ctx context.Context) error {
 	if w.Upstream == nil {
 		return fmt.Errorf("upstream manager is required")
 	}
+	pool := newMessageWorkPool(ctx, w, w.concurrency())
+	defer pool.wait()
 	for {
 		if err := w.ensureGroup(ctx); err != nil {
 			if ctx.Err() != nil {
@@ -87,7 +93,7 @@ func (w *Worker) Run(ctx context.Context) error {
 			sleep(ctx, 3*time.Second)
 			continue
 		}
-		if err := w.recoverPending(ctx); err != nil {
+		if err := w.recoverPending(ctx, pool); err != nil {
 			if ctx.Err() != nil {
 				return ctx.Err()
 			}
@@ -95,7 +101,7 @@ func (w *Worker) Run(ctx context.Context) error {
 			sleep(ctx, time.Second)
 			continue
 		}
-		if err := w.consumeOnce(ctx); err != nil {
+		if err := w.consumeOnce(ctx, pool); err != nil {
 			if ctx.Err() != nil {
 				return ctx.Err()
 			}
@@ -113,12 +119,16 @@ func (w *Worker) ensureGroup(ctx context.Context) error {
 	return err
 }

-func (w *Worker) consumeOnce(ctx context.Context) error {
+func (w *Worker) consumeOnce(ctx context.Context, pool *messageWorkPool) error {
+	available, err := pool.waitForCapacity(ctx)
+	if err != nil {
+		return err
+	}
 	streams, err := w.Redis.XReadGroup(ctx, &redis.XReadGroupArgs{
 		Group:    w.group(),
 		Consumer: w.consumer(),
 		Streams:  []string{w.stream(), ">"},
-		Count:    w.count(),
+		Count:    min(w.count(), int64(available)),
 		Block:    w.block(),
 	}).Result()
 	if errors.Is(err, redis.Nil) {
@@ -128,23 +138,29 @@ func (w *Worker) consumeOnce(ctx context.Context) error {
 		return err
 	}
 	for _, stream := range streams {
-		if err := w.processMessages(ctx, stream.Messages); err != nil {
-			return err
+		for _, message := range stream.Messages {
+			if !pool.dispatch(message) {
+				return fmt.Errorf("gateway submit worker capacity accounting mismatch")
+			}
 		}
 	}
 	return nil
 }

-func (w *Worker) recoverPending(ctx context.Context) error {
+func (w *Worker) recoverPending(ctx context.Context, pool *messageWorkPool) error {
 	start := "0-0"
 	for {
+		available := pool.available()
+		if available == 0 {
+			return nil
+		}
 		messages, next, err := w.Redis.XAutoClaim(ctx, &redis.XAutoClaimArgs{
 			Stream:   w.stream(),
 			Group:    w.group(),
 			Consumer: w.consumer(),
 			MinIdle:  w.minIdle(),
 			Start:    start,
-			Count:    w.count(),
+			Count:    min(w.count(), int64(available)),
 		}).Result()
 		if errors.Is(err, redis.Nil) {
 			return nil
@@ -156,8 +172,13 @@ func (w *Worker) recoverPending(ctx context.Context) error {
 			return nil
 		}
 		w.logf("gateway submit worker reclaimed %d pending message(s)", len(messages))
-		if err := w.processMessages(ctx, messages); err != nil {
-			return err
+		for _, message := range messages {
+			// An in-flight command can legitimately exceed MinIdle while waiting on a supplier.
+			// Rechecking both the local active set and Redis PEL closes the race where the
+			// original attempt ACKs between XAUTOCLAIM returning and local dispatch.
+			if err := pool.dispatchRecovered(ctx, message); err != nil {
+				return err
+			}
 		}
 		start = next
 		if next == "0-0" {
@@ -166,27 +187,111 @@ func (w *Worker) recoverPending(ctx context.Context) error {
 	}
 }

-func (w *Worker) processMessages(ctx context.Context, messages []redis.XMessage) error {
-	var group sync.WaitGroup
-	for _, message := range messages {
-		message := message
-		group.Add(1)
-		go func() {
-			defer group.Done()
-			if err := w.processMessage(ctx, message); err != nil {
-				w.logf("gateway submit worker message %s failed: %v", message.ID, err)
-			}
-		}()
+type messageWorkPool struct {
+	ctx       context.Context
+	worker    *Worker
+	slots     chan struct{}
+	completed chan struct{}
+	group     sync.WaitGroup
+	mu        sync.Mutex
+	active    map[string]struct{}
+}
+
+func newMessageWorkPool(ctx context.Context, worker *Worker, concurrency int) *messageWorkPool {
+	return &messageWorkPool{
+		ctx: ctx, worker: worker, slots: make(chan struct{}, concurrency),
+		completed: make(chan struct{}, concurrency), active: make(map[string]struct{}),
+	}
+}
+
+func (p *messageWorkPool) available() int {
+	return cap(p.slots) - len(p.slots)
+}
+
+func (p *messageWorkPool) waitForCapacity(ctx context.Context) (int, error) {
+	for p.available() == 0 {
+		select {
+		case <-ctx.Done():
+			return 0, ctx.Err()
+		case <-p.completed:
+		}
+	}
+	return p.available(), nil
+}
+
+func (p *messageWorkPool) dispatch(message redis.XMessage) bool {
+	p.mu.Lock()
+	if _, exists := p.active[message.ID]; exists {
+		p.mu.Unlock()
+		return true
+	}
+	select {
+	case p.slots <- struct{}{}:
+		p.active[message.ID] = struct{}{}
+		p.worker.inFlight.Add(1)
+		p.group.Add(1)
+		p.mu.Unlock()
+	case <-p.ctx.Done():
+		p.mu.Unlock()
+		return false
+	default:
+		p.mu.Unlock()
+		return false
+	}
+	go func() {
+		defer func() {
+			p.mu.Lock()
+			delete(p.active, message.ID)
+			p.mu.Unlock()
+			<-p.slots
+			p.worker.inFlight.Add(-1)
+			select {
+			case p.completed <- struct{}{}:
+			default:
+			}
+			p.group.Done()
+		}()
+		if err := p.worker.processMessage(p.ctx, message); err != nil {
+			p.worker.logf("gateway submit worker message %s failed: %v", message.ID, err)
+		}
+	}()
+	return true
+}
+
+func (p *messageWorkPool) dispatchRecovered(ctx context.Context, message redis.XMessage) error {
+	p.mu.Lock()
+	_, active := p.active[message.ID]
+	p.mu.Unlock()
+	if active {
+		return nil
+	}
+	pending, err := p.worker.Redis.XPendingExt(ctx, &redis.XPendingExtArgs{
+		Stream: p.worker.stream(), Group: p.worker.group(), Start: message.ID, End: message.ID, Count: 1,
+	}).Result()
+	if err != nil && !errors.Is(err, redis.Nil) {
+		return err
+	}
+	if len(pending) == 0 {
+		return nil
+	}
+	if !p.dispatch(message) {
+		return fmt.Errorf("gateway submit worker recovery capacity accounting mismatch")
 	}
-	group.Wait()
 	return nil
 }

+func (p *messageWorkPool) wait() {
+	p.group.Wait()
+}
+
 func (w *Worker) processMessage(ctx context.Context, message redis.XMessage) error {
 	command, err := CommandFromStreamValues(message.Values)
 	if err != nil {
 		return w.deadLetterMalformedMessage(ctx, message, err)
 	}
+	if !command.CreatedAt.IsZero() {
+		metrics.ObserveSubmitStage("stream_wait", true, time.Since(command.CreatedAt))
+	}
 	if err := w.handleCommand(ctx, command); err != nil {
 		if ctx.Err() != nil {
 			return ctx.Err()
@@ -208,11 +313,14 @@ func (w *Worker) processMessage(ctx context.Context, message redis.XMessage) err

 func (w *Worker) handleCommand(ctx context.Context, command queue.SubmitCommand) error {
 	startedAt := time.Now()
+	limitStartedAt := time.Now()
 	if w.Limiter != nil {
 		if _, err := w.Limiter.Wait(ctx, command.ChannelID, command.Route.RateLimitPerSecond); err != nil {
+			metrics.ObserveSubmitStage("rate_limit_wait", false, time.Since(limitStartedAt))
 			return err
 		}
 	}
+	metrics.ObserveSubmitStage("rate_limit_wait", true, time.Since(limitStartedAt))
 	submit := w.Submit
 	if submit == nil {
 		if w.Upstream == nil {
@@ -389,6 +497,21 @@ func (w *Worker) count() int64 {
 	return 10
 }

+func (w *Worker) concurrency() int {
+	if w.Concurrency > 0 {
+		return min(w.Concurrency, 1024)
+	}
+	return defaultConcurrency
+}
+
+func (w *Worker) ConfiguredConcurrency() int {
+	return w.concurrency()
+}
+
+func (w *Worker) InFlight() int64 {
+	return w.inFlight.Load()
+}
+
 func (w *Worker) minIdle() time.Duration {
 	if w.MinIdle > 0 {
 		return w.MinIdle
@@ -2,6 +2,7 @@ package submitworker

 import (
 	"context"
+	"sync/atomic"
 	"testing"
 	"time"

@@ -121,11 +122,21 @@ func TestMinIdleDefaultsToThirtySeconds(t *testing.T) {
 	}
 }

-func TestProcessMessagesDoesNotLetOneChannelBlockAnother(t *testing.T) {
+func TestConcurrencyUsesBoundedDefaultAndMaximum(t *testing.T) {
+	if got := (&Worker{}).ConfiguredConcurrency(); got != defaultConcurrency {
+		t.Fatalf("default concurrency = %d, want %d", got, defaultConcurrency)
+	}
+	if got := (&Worker{Concurrency: 2048}).ConfiguredConcurrency(); got != 1024 {
+		t.Fatalf("capped concurrency = %d, want 1024", got)
+	}
+}
+
+func TestMessageWorkPoolContinuouslyRefillsWithoutWaitingForSlowSibling(t *testing.T) {
 	mr := miniredis.RunT(t)
 	client := redis.NewClient(&redis.Options{Addr: mr.Addr()})
 	startedA := make(chan struct{})
 	startedB := make(chan struct{})
+	startedC := make(chan struct{})
 	releaseA := make(chan struct{})
 	worker := &Worker{
 		Redis: client,
@@ -136,19 +147,17 @@ func TestProcessMessagesDoesNotLetOneChannelBlockAnother(t *testing.T) {
 				<-releaseA
 			case "channel-b":
 				close(startedB)
+			case "channel-c":
+				close(startedC)
 			}
 			return queue.SubmitResult{SubmitStatus: "accepted"}, nil
 		},
 	}
-	messages := []redis.XMessage{
-		{ID: "1-0", Values: submitCommandValues("message-a", "channel-a")},
-		{ID: "2-0", Values: submitCommandValues("message-b", "channel-b")},
+	pool := newMessageWorkPool(context.Background(), worker, 2)
+	if !pool.dispatch(redis.XMessage{ID: "1-0", Values: submitCommandValues("message-a", "channel-a")}) ||
+		!pool.dispatch(redis.XMessage{ID: "2-0", Values: submitCommandValues("message-b", "channel-b")}) {
+		t.Fatal("initial messages were not dispatched")
 	}
-	done := make(chan struct{})
-	go func() {
-		_ = worker.processMessages(context.Background(), messages)
-		close(done)
-	}()
 	select {
 	case <-startedA:
 	case <-time.After(time.Second):
@@ -159,11 +168,125 @@ func TestProcessMessagesDoesNotLetOneChannelBlockAnother(t *testing.T) {
 	case <-time.After(200 * time.Millisecond):
 		t.Fatal("channel-b was blocked by channel-a")
 	}
-	close(releaseA)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if _, err := pool.waitForCapacity(ctx); err != nil {
+		t.Fatalf("wait for refill capacity: %v", err)
+	}
+	if !pool.dispatch(redis.XMessage{ID: "3-0", Values: submitCommandValues("message-c", "channel-c")}) {
+		t.Fatal("refill message was not dispatched")
+	}
 	select {
-	case <-done:
-	case <-time.After(time.Second):
-		t.Fatal("message batch did not complete")
+	case <-startedC:
+	case <-time.After(200 * time.Millisecond):
+		t.Fatal("pool waited for the slow sibling instead of refilling its free slot")
+	}
+	close(releaseA)
+	pool.wait()
+}
+
+func TestMessageWorkPoolAcknowledgesFastMessageBeforeSlowSiblingCompletes(t *testing.T) {
+	mr := miniredis.RunT(t)
+	client := redis.NewClient(&redis.Options{Addr: mr.Addr()})
+	worker := &Worker{Redis: client, Stream: "gateway.submit.commands", Group: "cmpp-gateway"}
+	ctx := context.Background()
+	if err := worker.ensureGroup(ctx); err != nil {
+		t.Fatalf("ensureGroup: %v", err)
+	}
+	for _, entry := range []struct{ id, messageID, channelID string }{
+		{"1-0", "message-slow", "channel-slow"},
+		{"2-0", "message-fast", "channel-fast"},
+	} {
+		if err := client.XAdd(ctx, &redis.XAddArgs{Stream: worker.stream(), ID: entry.id, Values: submitCommandValues(entry.messageID, entry.channelID)}).Err(); err != nil {
+			t.Fatalf("xadd %s: %v", entry.id, err)
+		}
+	}
+	streams, err := client.XReadGroup(ctx, &redis.XReadGroupArgs{Group: worker.group(), Consumer: worker.consumer(), Streams: []string{worker.stream(), ">"}, Count: 2}).Result()
+	if err != nil || len(streams) != 1 || len(streams[0].Messages) != 2 {
+		t.Fatalf("xreadgroup: streams=%+v err=%v", streams, err)
+	}
+	slowStarted := make(chan struct{})
+	fastReturned := make(chan struct{})
+	releaseSlow := make(chan struct{})
+	worker.Submit = func(_ context.Context, command queue.SubmitCommand) (queue.SubmitResult, error) {
+		if command.ChannelID == "channel-slow" {
+			close(slowStarted)
+			<-releaseSlow
+		} else {
+			close(fastReturned)
+		}
+		return queue.SubmitResult{SubmitStatus: "accepted"}, nil
+	}
+	pool := newMessageWorkPool(ctx, worker, 2)
+	for _, message := range streams[0].Messages {
+		if !pool.dispatch(message) {
+			t.Fatalf("message %s was not dispatched", message.ID)
+		}
+	}
+	<-slowStarted
+	<-fastReturned
+	deadline := time.Now().Add(time.Second)
+	for {
+		pending, pendingErr := client.XPending(ctx, worker.stream(), worker.group()).Result()
+		if pendingErr != nil {
+			t.Fatalf("xpending: %v", pendingErr)
+		}
+		if pending.Count == 1 {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("pending count = %d, want 1 while slow sibling is still running", pending.Count)
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	close(releaseSlow)
+	pool.wait()
+}
+
+func TestPendingRecoveryDoesNotDuplicateAnActiveOrAlreadyAcknowledgedMessage(t *testing.T) {
+	mr := miniredis.RunT(t)
+	client := redis.NewClient(&redis.Options{Addr: mr.Addr()})
+	worker := &Worker{Redis: client, Stream: "gateway.submit.commands", Group: "cmpp-gateway", MinIdle: time.Millisecond}
+	ctx := context.Background()
+	if err := worker.ensureGroup(ctx); err != nil {
+		t.Fatalf("ensureGroup: %v", err)
+	}
+	if err := client.XAdd(ctx, &redis.XAddArgs{Stream: worker.stream(), ID: "3-0", Values: submitCommandValues("message-active", "channel-active")}).Err(); err != nil {
+		t.Fatalf("xadd: %v", err)
+	}
+	streams, err := client.XReadGroup(ctx, &redis.XReadGroupArgs{Group: worker.group(), Consumer: worker.consumer(), Streams: []string{worker.stream(), ">"}, Count: 1}).Result()
+	if err != nil || len(streams) != 1 || len(streams[0].Messages) != 1 {
+		t.Fatalf("xreadgroup: streams=%+v err=%v", streams, err)
+	}
+	message := streams[0].Messages[0]
+	started := make(chan struct{})
+	release := make(chan struct{})
+	var submits atomic.Int32
+	worker.Submit = func(_ context.Context, _ queue.SubmitCommand) (queue.SubmitResult, error) {
+		submits.Add(1)
+		close(started)
+		<-release
+		return queue.SubmitResult{SubmitStatus: "accepted"}, nil
+	}
+	pool := newMessageWorkPool(ctx, worker, 2)
+	if !pool.dispatch(message) {
+		t.Fatal("active message was not dispatched")
+	}
+	<-started
+	time.Sleep(5 * time.Millisecond)
+	if err := worker.recoverPending(ctx, pool); err != nil {
+		t.Fatalf("recoverPending: %v", err)
+	}
+	if got := submits.Load(); got != 1 {
+		t.Fatalf("active message submit count = %d, want 1", got)
+	}
+	close(release)
+	pool.wait()
+	if err := pool.dispatchRecovered(ctx, message); err != nil {
+		t.Fatalf("dispatch acknowledged recovery: %v", err)
+	}
+	if got := submits.Load(); got != 1 {
+		t.Fatalf("acknowledged message submit count = %d, want 1", got)
 	}
 }