internal/mtprotoedge/inbound_rpc_test.go

Rama main · réplica de lectura del repositorio original.

package mtprotoedge

import (
	"context"
	"errors"
	"sync/atomic"
	"testing"
	"time"
)

func newInboundTestConn(s *inboundRPCScheduler, maxInflight, queueSize int, timeout time.Duration) *Conn {
	c := &Conn{metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s, maxInflight, queueSize, timeout)
	return c
}

func TestInboundRPCSchedulerIsLazyPerConnectionAndServer(t *testing.T) {
	scheduler := newInboundRPCScheduler(4, 16, 1<<20)
	scheduler.start()
	c := newInboundTestConn(scheduler, 2, 4, time.Second)
	defer func() {
		c.closeInboundRPCScheduler()
		scheduler.stop(time.Second)
	}()

	if c.rpcQueue != nil {
		t.Fatal("new connection eagerly allocated an inbound queue")
	}
	scheduler.lifecycleMu.Lock()
	workersStarted := scheduler.workersStarted
	scheduler.lifecycleMu.Unlock()
	if workersStarted {
		t.Fatal("empty server eagerly started inbound RPC workers")
	}
}

func TestInboundRPCSchedulerBoundsConcurrentWork(t *testing.T) {
	scheduler := newInboundRPCScheduler(2, 32, 1<<20)
	scheduler.start()
	c := newInboundTestConn(scheduler, 2, 4, time.Second)
	defer func() {
		c.closeInboundRPCScheduler()
		scheduler.stop(time.Second)
	}()

	var active atomic.Int64
	var maxActive atomic.Int64
	var done atomic.Int64
	started := make(chan struct{}, 6)
	release := make(chan struct{})
	task := inboundRPC{
		method: "test.method",
		run: func(ctx context.Context) error {
			cur := active.Add(1)
			for {
				old := maxActive.Load()
				if cur <= old || maxActive.CompareAndSwap(old, cur) {
					break
				}
			}
			started <- struct{}{}
			select {
			case <-release:
			case <-ctx.Done():
			}
			active.Add(-1)
			done.Add(1)
			return nil
		},
	}

	for i := 0; i < 2; i++ {
		if err := c.enqueueInboundRPC(context.Background(), task); err != nil {
			t.Fatalf("enqueue active task %d: %v", i, err)
		}
	}
	for i := 0; i < 2; i++ {
		select {
		case <-started:
		case <-time.After(time.Second):
			t.Fatal("timed out waiting for active rpc workers")
		}
	}
	for i := 0; i < 4; i++ {
		if err := c.enqueueInboundRPC(context.Background(), task); err != nil {
			t.Fatalf("enqueue queued task %d: %v", i, err)
		}
	}
	if err := c.enqueueInboundRPC(context.Background(), task); !errors.Is(err, ErrInboundRPCQueueFull) {
		t.Fatalf("enqueue over capacity err = %v, want ErrInboundRPCQueueFull", err)
	}
	if got := maxActive.Load(); got != 2 {
		t.Fatalf("max active = %d, want 2", got)
	}

	close(release)
	deadline := time.After(2 * time.Second)
	for done.Load() != 6 {
		select {
		case <-deadline:
			t.Fatalf("done = %d, want 6", done.Load())
		default:
			time.Sleep(10 * time.Millisecond)
		}
	}
	if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
		t.Fatalf("global budget after completion = (%d tasks, %d bytes), want zero", tasks, bytes)
	}
}

func TestInboundRPCSchedulerSkipsUnresolvedDependencyGate(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 8, 1<<20)
	scheduler.start()
	c := newInboundTestConn(scheduler, 1, 4, time.Second)
	defer func() {
		c.closeInboundRPCScheduler()
		scheduler.stop(time.Second)
	}()

	blockedRan := make(chan struct{}, 1)
	independentRan := make(chan struct{}, 1)
	gate := newInboundRPCGate(1, c.wakeInboundRPC)
	gate.resolve(true) // subscriber-installation sentinel; one dependency remains.
	if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
		method: "invokeAfter",
		gate:   gate,
		run: func(context.Context) error {
			blockedRan <- struct{}{}
			return nil
		},
	}); err != nil {
		t.Fatal(err)
	}
	if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
		method: "independent",
		run: func(context.Context) error {
			independentRan <- struct{}{}
			return nil
		},
	}); err != nil {
		t.Fatal(err)
	}

	select {
	case <-independentRan:
	case <-time.After(time.Second):
		t.Fatal("independent task was starved behind unresolved invokeAfter")
	}
	select {
	case <-blockedRan:
		t.Fatal("invokeAfter ran before its dependency completed")
	default:
	}
	gate.resolve(true)
	select {
	case <-blockedRan:
	case <-time.After(time.Second):
		t.Fatal("resolved invokeAfter was not rescheduled")
	}
}

func TestInboundRPCSchedulerFairAcrossConnections(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 16, 1<<20)
	c1 := newInboundTestConn(scheduler, 1, 4, time.Second)
	c2 := newInboundTestConn(scheduler, 1, 4, time.Second)
	defer func() {
		c1.closeInboundRPCScheduler()
		c2.closeInboundRPCScheduler()
		scheduler.stop(time.Second)
	}()

	order := make(chan string, 3)
	enqueue := func(c *Conn, label string) {
		t.Helper()
		if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
			method: label,
			run: func(context.Context) error {
				order <- label
				return nil
			},
		}); err != nil {
			t.Fatalf("enqueue %s: %v", label, err)
		}
	}

	// First form the [c1, c2] ready order before workers start. c1 runs only one
	// task at a time then returns to the tail, so c2 must get a chance to run before
	// c1's second task.
	enqueue(c1, "c1-first")
	enqueue(c1, "c1-second")
	enqueue(c2, "c2-first")
	scheduler.start()

	want := []string{"c1-first", "c2-first", "c1-second"}
	for i := range want {
		select {
		case got := <-order:
			if got != want[i] {
				t.Fatalf("execution[%d] = %q, want %q", i, got, want[i])
			}
		case <-time.After(time.Second):
			t.Fatalf("timed out waiting for execution[%d]", i)
		}
	}
}

func TestInboundRPCBudgetReservedBeforeCommitAndFullyReturned(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 2, 10)
	c1 := newInboundTestConn(scheduler, 1, 4, time.Second)
	c2 := newInboundTestConn(scheduler, 1, 4, time.Second)
	defer func() {
		c1.closeInboundRPCScheduler()
		c2.closeInboundRPCScheduler()
		scheduler.stop(time.Second)
	}()

	r1, err := c1.reserveInboundRPC(context.Background(), "one", 6)
	if err != nil {
		t.Fatalf("reserve first body: %v", err)
	}
	if tasks, bytes := scheduler.budgetSnapshot(); tasks != 1 || bytes != 6 {
		t.Fatalf("budget after first pre-Copy reservation = (%d, %d), want (1, 6)", tasks, bytes)
	}
	if _, err := c2.reserveInboundRPC(context.Background(), "too-large", 5); !errors.Is(err, ErrInboundRPCQueueFull) {
		t.Fatalf("reserve over byte budget err = %v, want queue full", err)
	}
	r2, err := c2.reserveInboundRPC(context.Background(), "two", 4)
	if err != nil {
		t.Fatalf("reserve second body: %v", err)
	}
	if _, err := c1.reserveInboundRPC(context.Background(), "too-many", 0); !errors.Is(err, ErrInboundRPCQueueFull) {
		t.Fatalf("reserve over task budget err = %v, want queue full", err)
	}

	r1.abort()
	r2.abort()
	if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
		t.Fatalf("budget after aborts = (%d, %d), want zero", tasks, bytes)
	}
	if got := c1.inflightRPCBytes.Load(); got != 0 {
		t.Fatalf("c1 inflight bytes = %d, want zero", got)
	}
	if got := c2.inflightRPCBytes.Load(); got != 0 {
		t.Fatalf("c2 inflight bytes = %d, want zero", got)
	}
}

func TestInboundRPCPerConnectionByteBudgetRejectedBeforeCommit(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 2, int64(maxInflightRPCBytes)+1)
	c := newInboundTestConn(scheduler, 1, 2, time.Second)
	defer func() {
		c.closeInboundRPCScheduler()
		scheduler.stop(time.Second)
	}()

	if _, err := c.reserveInboundRPC(context.Background(), "oversized", maxInflightRPCBytes+1); !errors.Is(err, ErrInboundRPCQueueFull) {
		t.Fatalf("reserve over per-connection byte budget err = %v, want queue full", err)
	}
	if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
		t.Fatalf("global budget after per-connection rejection = (%d, %d), want zero", tasks, bytes)
	}
	if got := c.inflightRPCBytes.Load(); got != 0 {
		t.Fatalf("connection bytes after rejection = %d, want zero", got)
	}
}

func TestInboundRPCCommitRacingCloseReturnsReservation(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 4, 1<<20)
	c := newInboundTestConn(scheduler, 1, 2, time.Second)
	defer scheduler.stop(time.Second)

	reservation, err := c.reserveInboundRPC(context.Background(), "closing", 13)
	if err != nil {
		t.Fatalf("reserve: %v", err)
	}
	closed := make(chan struct{})
	go func() {
		c.closeInboundRPCScheduler()
		close(closed)
	}()
	deadline := time.Now().Add(time.Second)
	for {
		c.rpcMu.Lock()
		isClosed := c.rpcClosed
		c.rpcMu.Unlock()
		if isClosed {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("connection scheduler was not marked closed")
		}
		time.Sleep(time.Millisecond)
	}

	if err := reservation.commit(inboundRPC{run: func(context.Context) error { return nil }}); !errors.Is(err, ErrConnClosed) {
		t.Fatalf("commit after close err = %v, want ErrConnClosed", err)
	}
	select {
	case <-closed:
	case <-time.After(time.Second):
		t.Fatal("close did not finish after reservation commit")
	}
	if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
		t.Fatalf("global budget after close/commit race = (%d, %d), want zero", tasks, bytes)
	}
	if got := c.inflightRPCBytes.Load(); got != 0 {
		t.Fatalf("connection bytes after close/commit race = %d, want zero", got)
	}
}

func TestInboundRPCSchedulerCloseRemovesReadyTokenBeforeStart(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 1, 1<<20)
	defer scheduler.stop(time.Second)

	// A bounded ready channel used to retain one stale token per closed connection. With workers
	// not started yet, the second connection then blocked forever trying to publish its token even
	// though the first connection had returned every task/byte budget.
	for i := 0; i < 32; i++ {
		c := newInboundTestConn(scheduler, 1, 1, time.Second)
		done := make(chan error, 1)
		go func() {
			done <- c.enqueueInboundRPC(context.Background(), inboundRPC{
				method: "close-before-start",
				run:    func(context.Context) error { return nil },
			})
		}()
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("enqueue iteration %d: %v", i, err)
			}
		case <-time.After(time.Second):
			t.Fatalf("enqueue iteration %d blocked behind a stale ready token", i)
		}
		c.closeInboundRPCScheduler()
		if got := scheduler.readyLen(); got != 0 {
			t.Fatalf("ready tokens after close iteration %d = %d, want zero", i, got)
		}
	}
	if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
		t.Fatalf("budget after close churn = (%d, %d), want zero", tasks, bytes)
	}
}

func TestInboundRPCExpiredInQueueNeverRunsAndSignalsTimeout(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 8, 1<<20)
	scheduler.start()
	c := newInboundTestConn(scheduler, 1, 4, 40*time.Millisecond)
	defer func() {
		c.closeInboundRPCScheduler()
		scheduler.stop(time.Second)
	}()

	started := make(chan struct{})
	release := make(chan struct{})
	if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
		method: "blocker",
		size:   7,
		run: func(context.Context) error {
			close(started)
			<-release // deliberately ignore the deadline, ensuring the next item expires in the queue.
			return nil
		},
	}); err != nil {
		t.Fatalf("enqueue blocker: %v", err)
	}
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("blocker did not start")
	}

	var ran atomic.Bool
	timedOut := make(chan struct{})
	if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
		method: "expires",
		size:   11,
		onTimeout: func() {
			close(timedOut)
		},
		run: func(context.Context) error {
			ran.Store(true)
			return nil
		},
	}); err != nil {
		t.Fatalf("enqueue expiring task: %v", err)
	}
	select {
	case <-timedOut:
	case <-time.After(time.Second):
		t.Fatal("queued task did not signal timeout while the worker was still blocked")
	}
	deadline := time.Now().Add(time.Second)
	for {
		tasks, bytes := scheduler.budgetSnapshot()
		if tasks == 1 && bytes == 7 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("budget while blocker still runs = (%d, %d), want only blocker (1, 7)", tasks, bytes)
		}
		time.Sleep(time.Millisecond)
	}
	close(release)
	if ran.Load() {
		t.Fatal("expired queued task entered business handler")
	}

	deadline = time.Now().Add(time.Second)
	for {
		tasks, bytes := scheduler.budgetSnapshot()
		if tasks == 0 && bytes == 0 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("budget after timeout = (%d, %d), want zero", tasks, bytes)
		}
		time.Sleep(time.Millisecond)
	}
}

func TestInboundRPCCloseDisarmsQueuedTimeout(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 8, 1<<20)
	c := newInboundTestConn(scheduler, 1, 4, 30*time.Millisecond)
	defer scheduler.stop(time.Second)

	timedOut := make(chan struct{}, 1)
	if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
		method: "queued",
		size:   11,
		onTimeout: func() {
			timedOut <- struct{}{}
		},
	}); err != nil {
		t.Fatalf("enqueue queued task: %v", err)
	}
	c.closeInboundRPCScheduler()
	time.Sleep(60 * time.Millisecond)
	select {
	case <-timedOut:
		t.Fatal("connection close emitted a queued RPC timeout")
	default:
	}
}

func TestInboundRPCRunningDeadlineCancelsWithoutEarlyTimeout(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 8, 1<<20)
	scheduler.start()
	c := newInboundTestConn(scheduler, 1, 4, 30*time.Millisecond)
	defer func() {
		c.closeInboundRPCScheduler()
		scheduler.stop(time.Second)
	}()

	started := make(chan struct{})
	release := make(chan struct{})
	timedOut := make(chan struct{}, 1)
	if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
		method: "running",
		size:   7,
		onTimeout: func() {
			timedOut <- struct{}{}
		},
		run: func(context.Context) error {
			close(started)
			<-release
			return nil
		},
	}); err != nil {
		t.Fatalf("enqueue running task: %v", err)
	}
	<-started
	time.Sleep(80 * time.Millisecond)
	select {
	case <-timedOut:
		t.Fatal("running task emitted an early timeout before handler convergence")
	default:
	}
	if tasks, bytes := scheduler.budgetSnapshot(); tasks != 1 || bytes != 7 {
		t.Fatalf("running body budget after timeout = (%d, %d), want retained (1, 7)", tasks, bytes)
	}
	close(release)
	deadline := time.Now().Add(time.Second)
	for {
		tasks, bytes := scheduler.budgetSnapshot()
		if tasks == 0 && bytes == 0 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("running body budget after completion = (%d, %d), want zero", tasks, bytes)
		}
		time.Sleep(time.Millisecond)
	}
}

func TestInboundRPCCloseDrainsQueueAndReturnsBudgets(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 8, 1<<20)
	scheduler.start()
	c := newInboundTestConn(scheduler, 1, 4, time.Second)
	defer scheduler.stop(time.Second)

	started := make(chan struct{})
	if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
		method: "running",
		size:   7,
		run: func(ctx context.Context) error {
			close(started)
			<-ctx.Done()
			return ctx.Err()
		},
	}); err != nil {
		t.Fatalf("enqueue running task: %v", err)
	}
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("running task did not start")
	}

	var queuedRan atomic.Bool
	if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
		method: "queued",
		size:   11,
		run: func(context.Context) error {
			queuedRan.Store(true)
			return nil
		},
	}); err != nil {
		t.Fatalf("enqueue queued task: %v", err)
	}
	if tasks, bytes := scheduler.budgetSnapshot(); tasks != 2 || bytes != 18 {
		t.Fatalf("budget before close = (%d, %d), want (2, 18)", tasks, bytes)
	}

	c.closeInboundRPCScheduler()
	if queuedRan.Load() {
		t.Fatal("queued task ran during connection close")
	}
	if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
		t.Fatalf("budget after close = (%d, %d), want zero", tasks, bytes)
	}
	if got := c.inflightRPCBytes.Load(); got != 0 {
		t.Fatalf("connection inflight bytes after close = %d, want zero", got)
	}
}

func TestInboundRPCReplayRestoreBarrierKeepsFollowingTaskOffWorkers(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 8, 1<<20)
	scheduler.start()
	c := newInboundTestConn(scheduler, 1, 4, time.Second)
	defer func() {
		c.closeInboundRPCScheduler()
		scheduler.stop(time.Second)
	}()

	finishFirst := c.beginRPCReplayRestore()
	finishSecond := c.beginRPCReplayRestore()
	ran := make(chan struct{})
	if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
		method: "following.naked.rpc",
		size:   4,
		run: func(context.Context) error {
			close(ran)
			return nil
		},
	}); err != nil {
		t.Fatalf("enqueue following task: %v", err)
	}

	select {
	case <-ran:
		t.Fatal("following RPC ran before replay restore completed")
	case <-time.After(30 * time.Millisecond):
	}
	finishFirst()
	select {
	case <-ran:
		t.Fatal("one of two replay restores released the scheduler early")
	case <-time.After(30 * time.Millisecond):
	}
	finishSecond()
	select {
	case <-ran:
	case <-time.After(time.Second):
		t.Fatal("following RPC did not run after the final replay restore")
	}
}