internal/mtprotoedge/outbound_required_control_test.go
Rama main · réplica de lectura del repositorio original.
package mtprotoedge
import (
"context"
"crypto/rand"
"errors"
"io"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
)
type gatedRequiredControlTransport struct {
started chan struct{}
release chan struct{}
sendErr error
startOnce sync.Once
closeOnce sync.Once
sends atomic.Int32
closes atomic.Int32
}
func newGatedRequiredControlTransport(sendErr error) *gatedRequiredControlTransport {
return &gatedRequiredControlTransport{
started: make(chan struct{}),
release: make(chan struct{}),
sendErr: sendErr,
}
}
func (t *gatedRequiredControlTransport) Send(context.Context, *bin.Buffer) error {
t.sends.Add(1)
t.startOnce.Do(func() { close(t.started) })
<-t.release
return t.sendErr
}
func (t *gatedRequiredControlTransport) Recv(context.Context, *bin.Buffer) error {
return io.EOF
}
func (t *gatedRequiredControlTransport) Close() error {
t.closes.Add(1)
t.closeOnce.Do(func() { close(t.release) })
return nil
}
func (t *gatedRequiredControlTransport) unblock() {
t.closeOnce.Do(func() { close(t.release) })
}
func TestServiceTaskResponsesWaitForPhysicalWrite(t *testing.T) {
tests := []struct {
name string
send func(*Server, context.Context, *Conn) error
}{
{
name: "pong",
send: func(s *Server, ctx context.Context, c *Conn) error {
return s.sendPong(ctx, c, 11, 22)
},
},
{
name: "future_salts",
send: func(s *Server, ctx context.Context, c *Conn) error {
return s.sendFutureSalts(ctx, c, 11, 32)
},
},
{
name: "msgs_state_info",
send: func(s *Server, ctx context.Context, c *Conn) error {
return s.sendMsgsStateInfo(ctx, c, 11, []byte{4})
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
tr := newGatedRequiredControlTransport(nil)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = newOutboundTrackedBudget(1 << 20)
srv := &Server{clock: clock.System}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
done := make(chan error, 1)
go func() {
done <- tc.send(srv, ctx, c)
}()
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("service response did not reach the physical writer")
}
select {
case err := <-done:
t.Fatalf("service response returned before physical write completed: %v", err)
case <-time.After(20 * time.Millisecond):
}
tr.unblock()
select {
case err := <-done:
if err != nil {
t.Fatalf("service response: %v", err)
}
case <-time.After(time.Second):
t.Fatal("service response did not return after physical write")
}
if c.isRetired() {
t.Fatal("successful service response terminally closed the connection")
}
if got := tr.sends.Load(); got != 1 {
t.Fatalf("physical sends = %d, want 1", got)
}
})
}
}
func TestSendRequiredControlWaitsForPhysicalWriteAndReturnsBudget(t *testing.T) {
tr := newGatedRequiredControlTransport(nil)
controlBudget := newOutboundTrackedBudget(1 << 20)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = controlBudget
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
done := make(chan error, 1)
go func() {
done <- c.SendRequiredControl(ctx, proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2})
}()
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("required control did not reach the physical writer")
}
select {
case err := <-done:
t.Fatalf("SendRequiredControl returned before physical write completed: %v", err)
case <-time.After(20 * time.Millisecond):
}
tr.unblock()
select {
case err := <-done:
if err != nil {
t.Fatalf("SendRequiredControl: %v", err)
}
case <-time.After(time.Second):
t.Fatal("SendRequiredControl did not return after physical write")
}
if c.isRetired() {
t.Fatal("successful required control terminally closed the connection")
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after non-pending physical write = %d, want 0", got)
}
if got := tr.sends.Load(); got != 1 {
t.Fatalf("physical sends = %d, want 1", got)
}
}
func TestSendRequiredControlReturnsAfterWriteWithoutWaitingForAck(t *testing.T) {
tr := &failAfterTransport{}
controlBudget := newOutboundTrackedBudget(1 << 20)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = controlBudget
created := &mt.NewSessionCreated{FirstMsgID: 1, UniqueID: 2, ServerSalt: 3}
encoded, err := encodeOutboundMessageWithoutSlot(created)
if err != nil {
t.Fatalf("encode new_session_created: %v", err)
}
if err := c.SendRequiredControl(context.Background(), proto.MessageFromServer, created); err != nil {
t.Fatalf("SendRequiredControl: %v", err)
}
if got := tr.stored.Load(); got != 1 {
t.Fatalf("completed physical sends = %d, want 1", got)
}
if got := controlBudget.snapshot(); got != int64(len(encoded.body)) {
t.Fatalf("pending control budget = %d, want %d until client ACK", got, len(encoded.body))
}
frame, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
if err != nil {
t.Fatalf("decrypt new_session_created: %v", err)
}
c.AckServerMessages([]int64{frame.MessageID})
deadline := time.Now().Add(time.Second)
for controlBudget.snapshot() != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after ACK = %d, want 0", got)
}
}
func TestSendRequiredControlQueueDeadlineTerminatesAndReturnsBudget(t *testing.T) {
tr := &failAfterTransport{}
controlBudget := newOutboundTrackedBudget(1 << 20)
c := &Conn{
transport: tr,
writer: tr,
metrics: NopMetrics{},
writeTimeout: time.Second,
outboundTrackedBudget: newOutboundTrackedBudget(1 << 20),
outboundControlTrackedBudget: controlBudget,
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
}
// No actor is running and the bounded control queue is full, so the parent
// deadline must cover queue admission and make the failure terminal.
c.outboundControl <- &outboundOp{kind: outboundAck}
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
defer cancel()
started := time.Now()
err := c.SendRequiredControl(ctx, proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2})
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("full control queue error = %v, want context deadline", err)
}
if elapsed := time.Since(started); elapsed > 500*time.Millisecond {
t.Fatalf("full control queue waited %v, want parent-deadline-bounded admission", elapsed)
}
if !c.isRetired() {
t.Fatal("required control queue failure did not terminally close the connection")
}
if got := tr.sends.Load(); got != 0 {
t.Fatalf("physical sends = %d, want 0", got)
}
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport closes = %d, want 1", got)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after queue timeout = %d, want 0", got)
}
}
func TestSendRequiredControlBlockedWriteUsesWholeOperationDeadline(t *testing.T) {
tr := newGatedRequiredControlTransport(io.ErrClosedPipe)
controlBudget := newOutboundTrackedBudget(1 << 20)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = controlBudget
c.writeTimeout = 25 * time.Millisecond
started := time.Now()
err := c.SendRequiredControl(context.Background(), proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2})
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("blocked required control error = %v, want context deadline", err)
}
if elapsed := time.Since(started); elapsed > 500*time.Millisecond {
t.Fatalf("blocked required control waited %v, want write-timeout-bounded operation", elapsed)
}
select {
case <-c.outboundDone:
case <-time.After(time.Second):
t.Fatal("outbound actor did not stop after required-control write timeout")
}
if !c.isRetired() {
t.Fatal("blocked required control did not terminally close the connection")
}
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport closes = %d, want 1", got)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after blocked write = %d, want 0", got)
}
}
func TestSendRequiredControlWriteFailureTerminatesAndReturnsBudget(t *testing.T) {
tr := &failAfterTransport{}
tr.failAt.Store(1)
controlBudget := newOutboundTrackedBudget(1 << 20)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = controlBudget
err := c.SendRequiredControl(context.Background(), proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2})
if err == nil {
t.Fatal("write-failed required control unexpectedly succeeded")
}
select {
case <-c.outboundDone:
case <-time.After(time.Second):
t.Fatal("outbound actor did not stop after required-control write failure")
}
if !c.isRetired() {
t.Fatal("write-failed required control did not terminally close the connection")
}
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport closes = %d, want 1", got)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after write failure = %d, want 0", got)
}
}
func TestSendRequiredControlBudgetFailureIsTerminal(t *testing.T) {
tr := &failAfterTransport{}
controlBudget := newOutboundTrackedBudget(1)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = controlBudget
err := c.SendRequiredControl(context.Background(), proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2})
if !errors.Is(err, ErrOutboundTrackedBudget) {
t.Fatalf("required control over budget = %v, want ErrOutboundTrackedBudget", err)
}
select {
case <-c.outboundDone:
case <-time.After(time.Second):
t.Fatal("outbound actor did not stop after required-control budget failure")
}
if !c.isRetired() {
t.Fatal("required-control budget failure did not terminally close the connection")
}
if got := tr.sends.Load(); got != 0 {
t.Fatalf("budget-rejected required control wrote %d frames, want 0", got)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after reservation failure = %d, want 0", got)
}
}