411 lines
13 KiB
Go
411 lines
13 KiB
Go
package mtprotoedge
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestInboundRPCBatchReservationRejectsGloballyWithoutPartialBudget(t *testing.T) {
|
|
scheduler := newInboundRPCScheduler(1, 2, 10)
|
|
c := newInboundTestConn(scheduler, 1, 8, time.Second)
|
|
defer func() {
|
|
c.closeInboundRPCScheduler()
|
|
scheduler.stop(time.Second)
|
|
}()
|
|
|
|
_, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
|
|
{method: "one", size: 3},
|
|
{method: "two", size: 3},
|
|
{method: "three", size: 3},
|
|
})
|
|
if !errors.Is(err, ErrInboundRPCQueueFull) {
|
|
t.Fatalf("reserve over global task budget err = %v, want ErrInboundRPCQueueFull", err)
|
|
}
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
|
t.Fatalf("global budget after atomic rejection = (%d, %d), want zero", tasks, bytes)
|
|
}
|
|
if got := c.inflightRPCBytes.Load(); got != 0 {
|
|
t.Fatalf("connection bytes after global rejection = %d, want zero", got)
|
|
}
|
|
c.rpcMu.Lock()
|
|
reserved := c.rpcReserved
|
|
queued := len(c.rpcQueue)
|
|
c.rpcMu.Unlock()
|
|
if reserved != 0 || queued != 0 {
|
|
t.Fatalf("connection state after global rejection = reserved %d queued %d, want zero", reserved, queued)
|
|
}
|
|
}
|
|
|
|
func TestInboundRPCBatchReservationRejectsConnectionWithoutLeakingGlobalBudget(t *testing.T) {
|
|
scheduler := newInboundRPCScheduler(1, 16, 1<<20)
|
|
c := newInboundTestConn(scheduler, 1, 2, time.Second)
|
|
defer func() {
|
|
c.closeInboundRPCScheduler()
|
|
scheduler.stop(time.Second)
|
|
}()
|
|
|
|
_, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
|
|
{method: "one", size: 3},
|
|
{method: "two", size: 5},
|
|
{method: "three", size: 7},
|
|
})
|
|
if !errors.Is(err, ErrInboundRPCQueueFull) {
|
|
t.Fatalf("reserve over connection queue budget err = %v, want ErrInboundRPCQueueFull", err)
|
|
}
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
|
t.Fatalf("global budget after connection rejection = (%d, %d), want zero", tasks, bytes)
|
|
}
|
|
if got := c.inflightRPCBytes.Load(); got != 0 {
|
|
t.Fatalf("connection bytes after connection rejection = %d, want zero", got)
|
|
}
|
|
c.rpcMu.Lock()
|
|
reserved := c.rpcReserved
|
|
queued := len(c.rpcQueue)
|
|
c.rpcMu.Unlock()
|
|
if reserved != 0 || queued != 0 {
|
|
t.Fatalf("connection state after connection rejection = reserved %d queued %d, want zero", reserved, queued)
|
|
}
|
|
}
|
|
|
|
func TestInboundRPCBatchReservationRejectsAggregateConnectionBytes(t *testing.T) {
|
|
scheduler := newInboundRPCScheduler(1, 8, int64(maxInflightRPCBytes)*2)
|
|
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
|
defer func() {
|
|
c.closeInboundRPCScheduler()
|
|
scheduler.stop(time.Second)
|
|
}()
|
|
|
|
halfPlusOne := maxInflightRPCBytes/2 + 1
|
|
_, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
|
|
{method: "one", size: halfPlusOne},
|
|
{method: "two", size: halfPlusOne},
|
|
})
|
|
if !errors.Is(err, ErrInboundRPCQueueFull) {
|
|
t.Fatalf("reserve over aggregate connection byte budget err = %v, want ErrInboundRPCQueueFull", err)
|
|
}
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
|
t.Fatalf("global budget after aggregate byte rejection = (%d, %d), want zero", tasks, bytes)
|
|
}
|
|
if got := c.inflightRPCBytes.Load(); got != 0 {
|
|
t.Fatalf("connection bytes after aggregate byte rejection = %d, want zero", got)
|
|
}
|
|
}
|
|
|
|
func TestInboundRPCBatchAbortReturnsEveryReservationExactlyOnce(t *testing.T) {
|
|
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
|
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
|
defer func() {
|
|
c.closeInboundRPCScheduler()
|
|
scheduler.stop(time.Second)
|
|
}()
|
|
|
|
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
|
|
{method: "one", size: 3},
|
|
{method: "two", size: 5},
|
|
{method: "three", size: 7},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("reserve batch: %v", err)
|
|
}
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 3 || bytes != 15 {
|
|
t.Fatalf("budget after reserve = (%d, %d), want (3, 15)", tasks, bytes)
|
|
}
|
|
reservation.abort()
|
|
reservation.abort()
|
|
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
|
t.Fatalf("budget after idempotent abort = (%d, %d), want zero", tasks, bytes)
|
|
}
|
|
if got := c.inflightRPCBytes.Load(); got != 0 {
|
|
t.Fatalf("connection bytes after abort = %d, want zero", got)
|
|
}
|
|
c.rpcMu.Lock()
|
|
reserved := c.rpcReserved
|
|
queued := len(c.rpcQueue)
|
|
c.rpcMu.Unlock()
|
|
if reserved != 0 || queued != 0 {
|
|
t.Fatalf("connection state after abort = reserved %d queued %d, want zero", reserved, queued)
|
|
}
|
|
}
|
|
|
|
func TestInboundRPCBatchReservationGrowTransfersBudgetAtomically(t *testing.T) {
|
|
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
|
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
|
defer func() {
|
|
c.closeInboundRPCScheduler()
|
|
scheduler.stop(time.Second)
|
|
}()
|
|
|
|
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
|
|
{method: "one", size: 3},
|
|
{method: "two", size: 5},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := reservation.growEntry(0, 11); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 2 || bytes != 16 {
|
|
t.Fatalf("grown global budget = %d/%d, want 2/16", tasks, bytes)
|
|
}
|
|
if got := c.inflightRPCBytes.Load(); got != 16 {
|
|
t.Fatalf("grown connection budget = %d, want 16", got)
|
|
}
|
|
if reservation.entries[0].size != 11 || reservation.totalSize != 16 {
|
|
t.Fatalf("grown reservation entry/total = %d/%d", reservation.entries[0].size, reservation.totalSize)
|
|
}
|
|
reservation.abort()
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
|
t.Fatalf("grown reservation abort leaked global budget %d/%d", tasks, bytes)
|
|
}
|
|
if got := c.inflightRPCBytes.Load(); got != 0 {
|
|
t.Fatalf("grown reservation abort leaked connection budget %d", got)
|
|
}
|
|
}
|
|
|
|
func TestInboundRPCBatchReservationGrowFailureKeepsOriginalBudget(t *testing.T) {
|
|
for _, test := range []struct {
|
|
name string
|
|
globalMax int64
|
|
targetSize int
|
|
}{
|
|
{name: "global", globalMax: 5, targetSize: 6},
|
|
{name: "connection", globalMax: int64(maxInflightRPCBytes) * 2, targetSize: maxInflightRPCBytes + 1},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
scheduler := newInboundRPCScheduler(1, 8, test.globalMax)
|
|
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
|
defer func() {
|
|
c.closeInboundRPCScheduler()
|
|
scheduler.stop(time.Second)
|
|
}()
|
|
|
|
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "one", size: 3}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := reservation.growEntry(0, test.targetSize); !errors.Is(err, ErrInboundRPCQueueFull) {
|
|
t.Fatalf("grow error = %v, want ErrInboundRPCQueueFull", err)
|
|
}
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 1 || bytes != 3 {
|
|
t.Fatalf("failed grow changed global budget %d/%d", tasks, bytes)
|
|
}
|
|
if got := c.inflightRPCBytes.Load(); got != 3 {
|
|
t.Fatalf("failed grow changed connection budget %d", got)
|
|
}
|
|
if reservation.entries[0].size != 3 || reservation.totalSize != 3 {
|
|
t.Fatalf("failed grow changed reservation entry/total = %d/%d", reservation.entries[0].size, reservation.totalSize)
|
|
}
|
|
reservation.abort()
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestInboundRPCBatchCommitAppendsAllAndSchedulesAtomically(t *testing.T) {
|
|
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
|
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
|
defer func() {
|
|
c.closeInboundRPCScheduler()
|
|
scheduler.stop(time.Second)
|
|
}()
|
|
|
|
specs := []inboundRPCSpec{
|
|
{method: "one", size: 3},
|
|
{method: "two", size: 5},
|
|
{method: "three", size: 7},
|
|
}
|
|
reservation, err := c.reserveInboundRPCBatch(context.Background(), specs)
|
|
if err != nil {
|
|
t.Fatalf("reserve batch: %v", err)
|
|
}
|
|
defer reservation.abort()
|
|
|
|
runs := make(chan string, len(specs))
|
|
tasks := make([]inboundRPC, len(specs))
|
|
for i, spec := range specs {
|
|
method := spec.method
|
|
tasks[i].run = func(context.Context) error {
|
|
runs <- method
|
|
return nil
|
|
}
|
|
}
|
|
if err := reservation.commit(tasks); err != nil {
|
|
t.Fatalf("commit batch: %v", err)
|
|
}
|
|
c.rpcMu.Lock()
|
|
queued := len(c.rpcQueue)
|
|
ready := c.rpcReady
|
|
c.rpcMu.Unlock()
|
|
if queued != len(specs) || !ready {
|
|
t.Fatalf("atomic queue state after commit = queued %d ready %v, want %d/true", queued, ready, len(specs))
|
|
}
|
|
if got := scheduler.readyLen(); got != 1 {
|
|
t.Fatalf("scheduler ready tokens after commit = %d, want one", got)
|
|
}
|
|
select {
|
|
case method := <-runs:
|
|
t.Fatalf("RPC %q ran before scheduler start", method)
|
|
default:
|
|
}
|
|
|
|
scheduler.start()
|
|
for _, want := range []string{"one", "two", "three"} {
|
|
select {
|
|
case got := <-runs:
|
|
if got != want {
|
|
t.Fatalf("execution order = %q, want %q", got, want)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatalf("timed out waiting for %q", want)
|
|
}
|
|
}
|
|
waitInboundRPCBatchBudget(t, scheduler, 0, 0)
|
|
}
|
|
|
|
func TestInboundRPCBatchCommitMismatchReleasesAllWithoutEnqueue(t *testing.T) {
|
|
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
|
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
|
defer func() {
|
|
c.closeInboundRPCScheduler()
|
|
scheduler.stop(time.Second)
|
|
}()
|
|
|
|
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
|
|
{method: "one", size: 3},
|
|
{method: "two", size: 5},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("reserve batch: %v", err)
|
|
}
|
|
if err := reservation.commit([]inboundRPC{{}}); !errors.Is(err, errInboundRPCBatchTaskCount) {
|
|
t.Fatalf("commit task mismatch err = %v, want %v", err, errInboundRPCBatchTaskCount)
|
|
}
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
|
t.Fatalf("budget after mismatched commit = (%d, %d), want zero", tasks, bytes)
|
|
}
|
|
if got := c.inflightRPCBytes.Load(); got != 0 {
|
|
t.Fatalf("connection bytes after mismatched commit = %d, want zero", got)
|
|
}
|
|
c.rpcMu.Lock()
|
|
reserved := c.rpcReserved
|
|
queued := len(c.rpcQueue)
|
|
c.rpcMu.Unlock()
|
|
if reserved != 0 || queued != 0 {
|
|
t.Fatalf("connection state after mismatched commit = reserved %d queued %d, want zero", reserved, queued)
|
|
}
|
|
}
|
|
|
|
func TestInboundRPCBatchCommitRacingCloseNeverPartiallyEnqueues(t *testing.T) {
|
|
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
|
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
|
defer scheduler.stop(time.Second)
|
|
|
|
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
|
|
{method: "one", size: 3},
|
|
{method: "two", size: 5},
|
|
{method: "three", size: 7},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("reserve batch: %v", err)
|
|
}
|
|
closed := make(chan struct{})
|
|
go func() {
|
|
c.closeInboundRPCScheduler()
|
|
close(closed)
|
|
}()
|
|
waitInboundRPCBatchConnClosed(t, c)
|
|
|
|
if err := reservation.commit(make([]inboundRPC, 3)); !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 batch commit returned its reservation")
|
|
}
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
|
t.Fatalf("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)
|
|
}
|
|
c.rpcMu.Lock()
|
|
queued := len(c.rpcQueue)
|
|
c.rpcMu.Unlock()
|
|
if queued != 0 {
|
|
t.Fatalf("queue after close/commit race = %d, want zero", queued)
|
|
}
|
|
}
|
|
|
|
func TestInboundRPCBatchCommitAfterTerminalFenceRejectsAll(t *testing.T) {
|
|
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
|
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
|
defer func() {
|
|
c.closeInboundRPCScheduler()
|
|
scheduler.stop(time.Second)
|
|
}()
|
|
|
|
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
|
|
{method: "one", size: 3},
|
|
{method: "two", size: 5},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("reserve batch: %v", err)
|
|
}
|
|
// Session replacement and revocation publish terminal before the slower
|
|
// physical-close path. A reservation held across that fence must not be able
|
|
// to append even one stale task.
|
|
c.retire()
|
|
if err := reservation.commit(make([]inboundRPC, 2)); !errors.Is(err, ErrConnClosed) {
|
|
t.Fatalf("commit after terminal fence err = %v, want ErrConnClosed", err)
|
|
}
|
|
|
|
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
|
t.Fatalf("budget after terminal commit rejection = (%d, %d), want zero", tasks, bytes)
|
|
}
|
|
if got := c.inflightRPCBytes.Load(); got != 0 {
|
|
t.Fatalf("connection bytes after terminal commit rejection = %d, want zero", got)
|
|
}
|
|
c.rpcMu.Lock()
|
|
reserved := c.rpcReserved
|
|
queued := len(c.rpcQueue)
|
|
c.rpcMu.Unlock()
|
|
if reserved != 0 || queued != 0 {
|
|
t.Fatalf("connection state after terminal commit rejection = reserved %d queued %d, want zero", reserved, queued)
|
|
}
|
|
}
|
|
|
|
func waitInboundRPCBatchBudget(t *testing.T, scheduler *inboundRPCScheduler, wantTasks int, wantBytes int64) {
|
|
t.Helper()
|
|
deadline := time.Now().Add(time.Second)
|
|
for {
|
|
tasks, bytes := scheduler.budgetSnapshot()
|
|
if tasks == wantTasks && bytes == wantBytes {
|
|
return
|
|
}
|
|
if time.Now().After(deadline) {
|
|
t.Fatalf("budget = (%d, %d), want (%d, %d)", tasks, bytes, wantTasks, wantBytes)
|
|
}
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
}
|
|
|
|
func waitInboundRPCBatchConnClosed(t *testing.T, c *Conn) {
|
|
t.Helper()
|
|
deadline := time.Now().Add(time.Second)
|
|
for {
|
|
c.rpcMu.Lock()
|
|
closed := c.rpcClosed
|
|
c.rpcMu.Unlock()
|
|
if closed {
|
|
return
|
|
}
|
|
if time.Now().After(deadline) {
|
|
t.Fatal("connection scheduler was not marked closed")
|
|
}
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
}
|