notify/transport_write_test.go

package notify

import (
	"b612.me/stario"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type serializedWriteTestConn struct {
	activeWrites int32
	concurrent   int32
	writeCount   int32
}

type deadlineRecordingWriteConn struct {
	mu        sync.Mutex
	deadlines []time.Time
}

func (c *deadlineRecordingWriteConn) Read([]byte) (int, error)        { return 0, net.ErrClosed }
func (c *deadlineRecordingWriteConn) Close() error                    { return nil }
func (c *deadlineRecordingWriteConn) LocalAddr() net.Addr             { return nil }
func (c *deadlineRecordingWriteConn) RemoteAddr() net.Addr            { return nil }
func (c *deadlineRecordingWriteConn) SetDeadline(time.Time) error     { return nil }
func (c *deadlineRecordingWriteConn) SetReadDeadline(time.Time) error { return nil }
func (c *deadlineRecordingWriteConn) Write(p []byte) (int, error)     { return len(p), nil }
func (c *deadlineRecordingWriteConn) SetWriteDeadline(deadline time.Time) error {
	c.mu.Lock()
	c.deadlines = append(c.deadlines, deadline)
	c.mu.Unlock()
	return nil
}

func (c *deadlineRecordingWriteConn) deadlineSnapshot() []time.Time {
	c.mu.Lock()
	defer c.mu.Unlock()
	return append([]time.Time(nil), c.deadlines...)
}

func assertDefaultBulkPhysicalDeadline(t *testing.T, deadlines []time.Time) {
	t.Helper()
	if len(deadlines) < 2 {
		t.Fatalf("physical write deadlines=%v, want bounded deadline followed by reset", deadlines)
	}
	remaining := time.Until(deadlines[0])
	if remaining < defaultBulkDataWriteTimeout-time.Second || remaining > defaultBulkDataWriteTimeout+time.Second {
		t.Fatalf("physical write deadline remaining=%v, want about %v", remaining, defaultBulkDataWriteTimeout)
	}
	if !deadlines[len(deadlines)-1].IsZero() {
		t.Fatalf("physical write deadline was not reset: %v", deadlines)
	}
}

type handoffWriteGateTestConn struct {
	firstStarted  chan struct{}
	secondStarted chan struct{}
	releaseFirst  chan struct{}
	writes        atomic.Int32
	active        atomic.Int32
	concurrent    atomic.Bool
}

func newHandoffWriteGateTestConn() *handoffWriteGateTestConn {
	return &handoffWriteGateTestConn{
		firstStarted:  make(chan struct{}),
		secondStarted: make(chan struct{}),
		releaseFirst:  make(chan struct{}),
	}
}

func (c *handoffWriteGateTestConn) Read([]byte) (int, error)         { return 0, net.ErrClosed }
func (c *handoffWriteGateTestConn) Close() error                     { return nil }
func (c *handoffWriteGateTestConn) LocalAddr() net.Addr              { return nil }
func (c *handoffWriteGateTestConn) RemoteAddr() net.Addr             { return nil }
func (c *handoffWriteGateTestConn) SetDeadline(time.Time) error      { return nil }
func (c *handoffWriteGateTestConn) SetReadDeadline(time.Time) error  { return nil }
func (c *handoffWriteGateTestConn) SetWriteDeadline(time.Time) error { return nil }
func (c *handoffWriteGateTestConn) Write(p []byte) (int, error) {
	active := c.active.Add(1)
	if active > 1 {
		c.concurrent.Store(true)
	}
	defer c.active.Add(-1)
	switch c.writes.Add(1) {
	case 1:
		close(c.firstStarted)
		<-c.releaseFirst
	case 2:
		close(c.secondStarted)
	}
	return len(p), nil
}

func (c *serializedWriteTestConn) Read([]byte) (int, error)         { return 0, net.ErrClosed }
func (c *serializedWriteTestConn) Close() error                     { return nil }
func (c *serializedWriteTestConn) LocalAddr() net.Addr              { return nil }
func (c *serializedWriteTestConn) RemoteAddr() net.Addr             { return nil }
func (c *serializedWriteTestConn) SetDeadline(time.Time) error      { return nil }
func (c *serializedWriteTestConn) SetReadDeadline(time.Time) error  { return nil }
func (c *serializedWriteTestConn) SetWriteDeadline(time.Time) error { return nil }

func (c *serializedWriteTestConn) Write(p []byte) (int, error) {
	if !atomic.CompareAndSwapInt32(&c.activeWrites, 0, 1) {
		atomic.StoreInt32(&c.concurrent, 1)
		return len(p), nil
	}
	time.Sleep(10 * time.Millisecond)
	atomic.AddInt32(&c.writeCount, 1)
	atomic.StoreInt32(&c.activeWrites, 0)
	return len(p), nil
}

func TestWriteFullToConnSerializesConcurrentWriters(t *testing.T) {
	conn := &serializedWriteTestConn{}
	payload := []byte("payload")

	var wg sync.WaitGroup
	for index := 0; index < 4; index++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := writeFullToConn(conn, payload); err != nil {
				t.Errorf("writeFullToConn failed: %v", err)
			}
		}()
	}
	wg.Wait()

	if atomic.LoadInt32(&conn.concurrent) != 0 {
		t.Fatal("detected concurrent conn.Write execution")
	}
	if got, want := atomic.LoadInt32(&conn.writeCount), int32(4); got != want {
		t.Fatalf("write count = %d, want %d", got, want)
	}
}

func TestTransportBindingAndRawHandoffSharePhysicalWriteGate(t *testing.T) {
	conn := newHandoffWriteGateTestConn()
	binding := newTransportBinding(conn, stario.NewQueue())
	bindingDone := make(chan error, 1)
	go func() {
		bindingDone <- binding.withConnWriteLock(func(conn net.Conn) error {
			return writeFullToConnUnlocked(conn, []byte("old-binding-frame"))
		})
	}()
	select {
	case <-conn.firstStarted:
	case <-time.After(time.Second):
		t.Fatal("binding write did not start")
	}

	rawDone := make(chan error, 1)
	go func() {
		rawDone <- withRawConnWriteLock(conn, func(conn net.Conn) error {
			return writeFullToConnUnlocked(conn, []byte("handoff-reply"))
		})
	}()
	select {
	case <-conn.secondStarted:
		close(conn.releaseFirst)
		<-bindingDone
		<-rawDone
		t.Fatal("raw handoff write entered Conn.Write before the binding frame completed")
	case <-time.After(50 * time.Millisecond):
	}
	close(conn.releaseFirst)
	for name, done := range map[string]<-chan error{"binding": bindingDone, "raw": rawDone} {
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("%s write failed: %v", name, err)
			}
		case <-time.After(time.Second):
			t.Fatalf("%s write did not finish", name)
		}
	}
	if conn.concurrent.Load() {
		t.Fatal("binding and raw handoff writes reached Conn.Write concurrently")
	}
}

func TestReleasedTransportBindingKeepsGateUntilActiveWriteFinishes(t *testing.T) {
	conn := newHandoffWriteGateTestConn()
	oldBinding := newTransportBinding(conn, stario.NewQueue())
	oldDone := make(chan error, 1)
	go func() {
		oldDone <- oldBinding.withConnWriteLock(func(conn net.Conn) error {
			return writeFullToConnUnlocked(conn, []byte("old-frame"))
		})
	}()
	select {
	case <-conn.firstStarted:
	case <-time.After(time.Second):
		t.Fatal("old binding write did not start")
	}

	stopDone := make(chan struct{})
	go func() {
		oldBinding.stopBackgroundWorkers()
		close(stopDone)
	}()
	deadline := time.Now().Add(time.Second)
	for !oldBinding.writeStopping.Load() && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if !oldBinding.writeStopping.Load() {
		close(conn.releaseFirst)
		<-oldDone
		<-stopDone
		t.Fatal("old binding did not enter write shutdown")
	}

	newBinding := newTransportBinding(conn, stario.NewQueue())
	newDone := make(chan error, 1)
	go func() {
		newDone <- newBinding.withConnWriteLock(func(conn net.Conn) error {
			return writeFullToConnUnlocked(conn, []byte("new-frame"))
		})
	}()
	select {
	case <-conn.secondStarted:
		close(conn.releaseFirst)
		<-oldDone
		<-newDone
		<-stopDone
		t.Fatal("replacement binding entered Conn.Write while the released binding still had an active write")
	case <-time.After(50 * time.Millisecond):
	}

	close(conn.releaseFirst)
	for name, done := range map[string]<-chan error{"old binding": oldDone, "new binding": newDone} {
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("%s write failed: %v", name, err)
			}
		case <-time.After(time.Second):
			t.Fatalf("%s write did not finish", name)
		}
	}
	select {
	case <-stopDone:
	case <-time.After(time.Second):
		t.Fatal("old binding shutdown did not finish")
	}
	newBinding.stopBackgroundWorkers()
	if conn.concurrent.Load() {
		t.Fatal("released and replacement bindings reached Conn.Write concurrently")
	}
}

func TestRawHandoffWriteDeadlineBoundsBindingGateWait(t *testing.T) {
	conn := &serializedWriteTestConn{}
	binding := newTransportBinding(conn, stario.NewQueue())
	locked := make(chan struct{})
	release := make(chan struct{})
	bindingDone := make(chan error, 1)
	go func() {
		bindingDone <- binding.withConnWriteLock(func(net.Conn) error {
			close(locked)
			<-release
			return nil
		})
	}()
	select {
	case <-locked:
	case <-time.After(time.Second):
		t.Fatal("binding write gate was not acquired")
	}
	callbackCalled := false
	err := withRawConnWriteLockDeadline(conn, time.Now().Add(30*time.Millisecond), func(net.Conn) error {
		callbackCalled = true
		return nil
	})
	close(release)
	if bindingErr := <-bindingDone; bindingErr != nil {
		t.Fatalf("binding write failed: %v", bindingErr)
	}
	if !isTimeoutLikeError(err) {
		t.Fatalf("raw handoff gate wait error=%v, want timeout", err)
	}
	if callbackCalled {
		t.Fatal("raw handoff callback ran without acquiring the shared binding gate")
	}
}

func newTestBulkBatchSender(binding *transportBinding) *bulkBatchSender {
	return newTestBulkBatchSenderWithWriteTimeout(binding, nil)
}

func newTestBulkBatchSenderWithWriteTimeout(binding *transportBinding, writeTimeout func() time.Duration) *bulkBatchSender {
	return newBulkBatchSender(binding, bulkBatchCodec{
		encodeSingle: func(frame bulkFastFrame) ([]byte, func(), error) {
			return append([]byte(nil), frame.Payload...), nil, nil
		},
		encodeBatch: func(frames []bulkFastFrame) ([]byte, func(), error) {
			payload, err := encodeBulkFastBatchPlain(frames)
			return payload, nil, err
		},
	}, writeTimeout)
}

func newTestStreamBatchSender(binding *transportBinding, writeTimeout func() time.Duration) *streamBatchSender {
	return newStreamBatchSender(binding, streamBatchCodec{
		encodeSingle: func(frame streamFastDataFrame) ([]byte, error) {
			return encodeStreamFastFramePayload(frame)
		},
		encodeBatch: func(frames []streamFastDataFrame) ([]byte, error) {
			return encodeStreamFastBatchPlain(frames)
		},
	}, writeTimeout)
}

func TestBulkBatchSenderRespectsWriteDeadlineWhenReceiverStalls(t *testing.T) {
	left, right := net.Pipe()
	defer left.Close()
	defer right.Close()

	sender := newTestBulkBatchSenderWithWriteTimeout(newTransportBinding(left, stario.NewQueue()), func() time.Duration {
		return 50 * time.Millisecond
	})

	errCh := make(chan error, 1)
	go func() {
		errCh <- sender.submitData(context.Background(), 1, 1, bulkFastPathVersionV1, []byte("payload"))
	}()

	select {
	case err := <-errCh:
		if err == nil {
			t.Fatal("sender.submit should fail when receiver stalls")
		}
		if !isTimeoutLikeError(err) {
			t.Fatalf("sender.submit error = %v, want timeout-like error", err)
		}
	case <-time.After(time.Second):
		t.Fatal("sender.submit should not hang when receiver stalls")
	}
}

func TestBulkBatchSenderUsesDefaultPhysicalWriteDeadline(t *testing.T) {
	conn := &deadlineRecordingWriteConn{}
	sender := newTestBulkBatchSender(newTransportBinding(conn, stario.NewQueue()))
	defer sender.stop()

	if err := sender.submitData(context.Background(), 1, 1, bulkFastPathVersionV1, []byte("payload")); err != nil {
		t.Fatalf("sender.submitData returned error: %v", err)
	}
	assertDefaultBulkPhysicalDeadline(t, conn.deadlineSnapshot())
}

func TestDedicatedBulkUsesDefaultPhysicalWriteDeadline(t *testing.T) {
	conn := &deadlineRecordingWriteConn{}
	if err := writeBulkDedicatedRecordWithDeadline(conn, []byte("payload"), time.Time{}); err != nil {
		t.Fatalf("writeBulkDedicatedRecordWithDeadline returned error: %v", err)
	}
	assertDefaultBulkPhysicalDeadline(t, conn.deadlineSnapshot())
}

func TestStreamBatchSenderRespectsBindingWriteDeadlineWhenReceiverStalls(t *testing.T) {
	left, right := net.Pipe()
	defer left.Close()
	defer right.Close()

	sender := newTestStreamBatchSender(newTransportBinding(left, stario.NewQueue()), func() time.Duration {
		return 50 * time.Millisecond
	})

	errCh := make(chan error, 1)
	go func() {
		errCh <- sender.submitData(context.Background(), 1, 1, streamFastPathVersionV2, []byte("payload"))
	}()

	select {
	case err := <-errCh:
		if err == nil {
			t.Fatal("sender.submit should fail when receiver stalls")
		}
		if !isTimeoutLikeError(err) {
			t.Fatalf("sender.submit error = %v, want timeout-like error", err)
		}
	case <-time.After(time.Second):
		t.Fatal("sender.submit should not hang when receiver stalls")
	}
}

func TestControlBatchSenderCarriesContextDeadlineIntoPhysicalWrite(t *testing.T) {
	left, right := net.Pipe()
	defer left.Close()
	defer right.Close()

	sender := newControlBatchSender(newTransportBinding(left, stario.NewQueue()))
	defer sender.stop()
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	err := sender.submitContext(ctx, []byte("deadline-control"), 0, controlPriorityNormal)
	if err == nil || !isTimeoutLikeError(err) {
		t.Fatalf("control submit error=%v, want timeout-like error", err)
	}
}

func TestBulkBatchSenderCarriesContextDeadlineIntoPhysicalWrite(t *testing.T) {
	left, right := net.Pipe()
	defer left.Close()
	defer right.Close()

	sender := newTestBulkBatchSender(newTransportBinding(left, stario.NewQueue()))
	defer sender.stop()
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	err := sender.submitData(ctx, 1, 1, bulkFastPathVersionV1, []byte("deadline-bulk"))
	if err == nil || !isTimeoutLikeError(err) {
		t.Fatalf("bulk submit error=%v, want timeout-like error", err)
	}
}

func TestStreamBatchSenderCarriesContextDeadlineIntoPhysicalWrite(t *testing.T) {
	left, right := net.Pipe()
	defer left.Close()
	defer right.Close()

	sender := newTestStreamBatchSender(newTransportBinding(left, stario.NewQueue()), nil)
	defer sender.stop()
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	err := sender.submitData(ctx, 1, 1, streamFastPathVersionV1, []byte("deadline-stream"))
	if err == nil || !isTimeoutLikeError(err) {
		t.Fatalf("stream submit error=%v, want timeout-like error", err)
	}
}

func TestTransportBindingStopWithCloseInterruptsPhysicalWrite(t *testing.T) {
	left, right := net.Pipe()
	defer right.Close()

	binding := newTransportBinding(left, stario.NewQueue())
	sender := newTestBulkBatchSender(binding)
	errCh := make(chan error, 1)
	go func() {
		errCh <- sender.submitData(context.Background(), 1, 1, bulkFastPathVersionV1, []byte("blocked"))
	}()
	time.Sleep(20 * time.Millisecond)
	stopCh := make(chan struct{})
	go func() {
		binding.stopBackgroundWorkersWithClose(true)
		close(stopCh)
	}()
	select {
	case <-stopCh:
	case <-time.After(time.Second):
		t.Fatal("stopping transport workers remained blocked on physical write")
	}
	select {
	case <-errCh:
	case <-time.After(time.Second):
		t.Fatal("blocked submit did not complete after transport close")
	}
}

func TestSameSocketHandoffBoundsBlockedSenderAndClosesSocket(t *testing.T) {
	left, right := net.Pipe()
	defer right.Close()

	binding := newTransportBinding(left, stario.NewQueue())
	sender := newTestBulkBatchSender(binding)
	binding.bulkMu.Lock()
	binding.bulkSender = sender
	binding.bulkMu.Unlock()
	writeDone := make(chan error, 1)
	go func() {
		writeDone <- sender.submitData(context.Background(), 1, 1, bulkFastPathVersionV1, []byte("handoff-blocked"))
	}()
	// Give the sender time to leave the queue and enter the physical write.
	time.Sleep(20 * time.Millisecond)

	stopDone := make(chan struct{})
	started := time.Now()
	go func() {
		// The next binding deliberately uses the same socket. The old sender is
		// already inside Conn.Write, so retaining the socket is unsafe.
		stopReplacedTransportBinding(binding, binding, false)
		close(stopDone)
	}()
	select {
	case <-stopDone:
		if elapsed := time.Since(started); elapsed > 3*transportWorkerStopWait {
			t.Fatalf("same-socket handoff stop took %v, want bounded", elapsed)
		}
	case <-time.After(3 * transportWorkerStopWait):
		t.Fatal("same-socket handoff remained blocked on an in-flight physical write")
	}
	select {
	case err := <-writeDone:
		if err == nil {
			t.Fatal("blocked sender unexpectedly succeeded after handoff forced socket close")
		}
	case <-time.After(time.Second):
		t.Fatal("blocked sender did not finish after handoff forced socket close")
	}
}

func TestClientCanceledPacketSendDoesNotWrite(t *testing.T) {
	conn := newBlockingPacketWriteConn()
	client := NewClient().(*ClientCommon)
	stopCtx, stopFn := context.WithCancel(context.Background())
	defer stopFn()
	client.setClientSessionRuntime(newClientSessionRuntime(conn, stopCtx, stopFn, stario.NewQueue(), 1))
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if err := client.writeControlPayloadToTransport(ctx, []byte("canceled-packet"), controlPriorityNormal); !errors.Is(err, context.Canceled) {
		t.Fatalf("canceled packet send error=%v, want context canceled", err)
	}
	if got := conn.writeCount.Load(); got != 0 {
		t.Fatalf("canceled packet send wrote %d packets", got)
	}
}

func TestServerCanceledUDPSendDoesNotWrite(t *testing.T) {
	server := NewServer().(*ServerCommon)
	server.alive.Store(true)
	stopCtx, stopFn := context.WithCancel(context.Background())
	defer stopFn()
	sender, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	if err != nil {
		t.Fatalf("ListenUDP sender failed: %v", err)
	}
	defer sender.Close()
	receiver, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	if err != nil {
		t.Fatalf("ListenUDP receiver failed: %v", err)
	}
	defer receiver.Close()
	server.setServerSessionRuntime(&serverSessionRuntime{stopCtx: stopCtx, stopFn: stopFn, udpListener: sender})
	logical := server.bootstrapAcceptedLogical("canceled-udp", receiver.LocalAddr(), nil)
	if logical == nil || logical.CurrentTransportConn() == nil {
		t.Fatal("failed to create UDP transport route")
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := server.sendTransportContext(ctx, logical.CurrentTransportConn(), TransferMsg{Key: "canceled", Type: MSG_ASYNC}); !errors.Is(err, context.Canceled) {
		t.Fatalf("canceled UDP send error=%v, want context canceled", err)
	}
	_ = receiver.SetReadDeadline(time.Now().Add(50 * time.Millisecond))
	var buf [256]byte
	if _, _, err := receiver.ReadFromUDP(buf[:]); err == nil {
		t.Fatal("canceled UDP send produced a datagram")
	}
}

func TestServerUDPWriteLockSerializesWriters(t *testing.T) {
	server := NewServer().(*ServerCommon)
	var active atomic.Int32
	var maxActive atomic.Int32
	var first sync.Once
	started := make(chan struct{})
	continueFirst := make(chan struct{})
	write := func() error {
		current := active.Add(1)
		for {
			previous := maxActive.Load()
			if current <= previous || maxActive.CompareAndSwap(previous, current) {
				break
			}
		}
		first.Do(func() {
			close(started)
			<-continueFirst
		})
		active.Add(-1)
		return nil
	}
	firstDone := make(chan error, 1)
	secondDone := make(chan error, 1)
	go func() { firstDone <- server.withUDPWriteLock(context.Background(), write) }()
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("first UDP writer did not start")
	}
	go func() { secondDone <- server.withUDPWriteLock(context.Background(), write) }()
	time.Sleep(30 * time.Millisecond)
	if got := maxActive.Load(); got != 1 {
		t.Fatalf("concurrent UDP writers observed = %d, want 1", got)
	}
	close(continueFirst)
	for _, done := range []<-chan error{firstDone, secondDone} {
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("UDP writer returned error: %v", err)
			}
		case <-time.After(time.Second):
			t.Fatal("UDP writer did not finish")
		}
	}
}

func TestServerTransportContextSendRejectsMissingRouteBeforeDereference(t *testing.T) {
	server := NewServer().(*ServerCommon)
	server.alive.Store(true)
	msg := TransferMsg{Key: "missing-route", Type: MSG_ASYNC}

	if _, err := server.sendTUTransportContext(context.Background(), nil, msg); !errors.Is(err, errTransportDetached) {
		t.Fatalf("missing TU transport error=%v, want transport detached", err)
	}
	if _, err := server.sendUDPTransportContext(context.Background(), nil, msg); !errors.Is(err, errTransportDetached) {
		t.Fatalf("missing UDP transport error=%v, want transport detached", err)
	}
}

func TestBulkBatchSenderFlushAggregatesAdaptivePayloadObservation(t *testing.T) {
	conn := &delayedWriteConn{delay: 20 * time.Millisecond}
	binding := newTransportBinding(conn, stario.NewQueue())
	sender := newTestBulkBatchSender(binding)
	payloadA := bytes.Repeat([]byte("a"), 128*1024)
	payloadB := bytes.Repeat([]byte("b"), 128*1024)

	err := sender.flush([]bulkBatchRequest{
		{
			ctx: context.Background(),
			frames: []bulkFastFrame{{
				Type:    bulkFastPayloadTypeData,
				DataID:  1,
				Seq:     1,
				Payload: payloadA,
			}},
			fastPathVersion: bulkFastPathVersionV1,
		},
		{
			ctx: context.Background(),
			frames: []bulkFastFrame{{
				Type:    bulkFastPayloadTypeData,
				DataID:  2,
				Seq:     1,
				Payload: payloadB,
			}},
			fastPathVersion: bulkFastPathVersionV1,
		},
	})
	if err != nil {
		t.Fatalf("flush failed: %v", err)
	}
	if got, want := binding.bulkAdaptiveSoftPayloadBytesSnapshot(), bulkAdaptiveSoftPayloadMinBytes; got != want {
		t.Fatalf("adaptive bulk soft payload = %d, want %d", got, want)
	}
}

func TestControlBatchSenderRespectsWriteDeadlineWhenReceiverStalls(t *testing.T) {
	left, right := net.Pipe()
	defer left.Close()
	defer right.Close()

	sender := newControlBatchSender(newTransportBinding(left, stario.NewQueue()))
	writeTimeout := 50 * time.Millisecond

	errCh := make(chan error, 1)
	go func() {
		errCh <- sender.submit([]byte("payload"), writeTimeout)
	}()

	select {
	case err := <-errCh:
		if err == nil {
			t.Fatal("sender.submit should fail when receiver stalls")
		}
		if !isTimeoutLikeError(err) {
			t.Fatalf("sender.submit error = %v, want timeout-like error", err)
		}
	case <-time.After(time.Second):
		t.Fatal("sender.submit should not hang when receiver stalls")
	}
}

func TestControlBatchDirectSubmitDoesNotRequireQueuedState(t *testing.T) {
	conn := &serializedWriteTestConn{}
	sender := newControlBatchSender(newTransportBinding(conn, stario.NewQueue()))
	defer sender.stop()

	submitted, err := sender.tryDirectSubmit(controlBatchRequest{
		ctx:       context.Background(),
		payload:   []byte("direct-control"),
		priority:  controlPriorityNormal,
		queueSize: int64(len("direct-control")),
	})
	if err != nil {
		t.Fatalf("tryDirectSubmit returned error: %v", err)
	}
	if !submitted {
		t.Fatal("tryDirectSubmit did not use the idle direct path")
	}
	if got := atomic.LoadInt32(&conn.writeCount); got == 0 {
		t.Fatal("direct submit did not write a framed payload")
	}
	if atomic.LoadInt32(&conn.concurrent) != 0 {
		t.Fatal("direct submit performed concurrent transport writes")
	}
}

func TestControlBatchQueuedSubmitAllocatesCompletionState(t *testing.T) {
	sender := &controlBatchSender{
		normalCh:   make(chan controlBatchRequest, 1),
		criticalCh: make(chan controlBatchRequest, 1),
		stopCh:     make(chan struct{}),
		doneCh:     make(chan struct{}),
		budgetWake: make(chan struct{}, 1),
	}
	sender.flushMu.Lock()
	defer sender.flushMu.Unlock()

	errCh := make(chan error, 1)
	go func() {
		errCh <- sender.submitContext(context.Background(), []byte("queued-control"), 0, controlPriorityNormal)
	}()

	var req controlBatchRequest
	select {
	case req = <-sender.normalCh:
	case <-time.After(time.Second):
		t.Fatal("control request did not enter the queue")
	}
	if req.done == nil {
		t.Fatal("queued control request has no completion channel")
	}
	if req.state == nil {
		t.Fatal("queued control request has no cancellation state")
	}
	sender.finishRequest(req, nil)

	select {
	case err := <-errCh:
		if err != nil {
			t.Fatalf("queued submit returned error: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("queued submit did not return after completion")
	}
}

type blockingRecordingControlConn struct {
	startOnce   sync.Once
	unblockOnce sync.Once
	startCh     chan struct{}
	unblockCh   chan struct{}

	mu   sync.Mutex
	data bytes.Buffer
}

func newBlockingRecordingControlConn() *blockingRecordingControlConn {
	return &blockingRecordingControlConn{
		startCh:   make(chan struct{}),
		unblockCh: make(chan struct{}),
	}
}

func (c *blockingRecordingControlConn) Read([]byte) (int, error) { return 0, net.ErrClosed }
func (c *blockingRecordingControlConn) Close() error             { return nil }
func (c *blockingRecordingControlConn) LocalAddr() net.Addr      { return nil }
func (c *blockingRecordingControlConn) RemoteAddr() net.Addr     { return nil }
func (c *blockingRecordingControlConn) SetDeadline(time.Time) error {
	return nil
}
func (c *blockingRecordingControlConn) SetReadDeadline(time.Time) error {
	return nil
}
func (c *blockingRecordingControlConn) SetWriteDeadline(time.Time) error {
	return nil
}

func (c *blockingRecordingControlConn) Write(p []byte) (int, error) {
	c.startOnce.Do(func() {
		close(c.startCh)
		<-c.unblockCh
	})
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.data.Write(p)
}

func (c *blockingRecordingControlConn) bytes() []byte {
	c.mu.Lock()
	defer c.mu.Unlock()
	return append([]byte(nil), c.data.Bytes()...)
}

func (c *blockingRecordingControlConn) unblock() {
	c.unblockOnce.Do(func() {
		close(c.unblockCh)
	})
}

func TestControlBatchSenderPrioritizesCriticalRequestAfterCurrentFlush(t *testing.T) {
	conn := newBlockingRecordingControlConn()
	sender := newControlBatchSender(newTransportBinding(conn, stario.NewQueue()))
	defer sender.stop()
	defer conn.unblock()

	firstDone := make(chan error, 1)
	go func() {
		firstDone <- sender.submitContext(context.Background(), []byte("control-first"), 0, controlPriorityNormal)
	}()
	select {
	case <-conn.startCh:
	case <-time.After(time.Second):
		t.Fatal("first control write did not start")
	}

	normalDone := make(chan error, 1)
	go func() {
		normalDone <- sender.submitContext(context.Background(), []byte("control-normal"), 0, controlPriorityNormal)
	}()
	criticalDone := make(chan error, 1)
	go func() {
		criticalDone <- sender.submitContext(context.Background(), []byte("control-critical"), 0, controlPriorityCritical)
	}()
	time.Sleep(20 * time.Millisecond)
	conn.unblock()

	for name, done := range map[string]<-chan error{
		"first": firstDone, "normal": normalDone, "critical": criticalDone,
	} {
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("%s submit failed: %v", name, err)
			}
		case <-time.After(time.Second):
			t.Fatalf("%s submit did not finish", name)
		}
	}

	written := conn.bytes()
	criticalAt := bytes.Index(written, []byte("control-critical"))
	normalAt := bytes.Index(written, []byte("control-normal"))
	if criticalAt < 0 || normalAt < 0 {
		t.Fatalf("missing control payloads in write: critical=%d normal=%d", criticalAt, normalAt)
	}
	if criticalAt > normalAt {
		t.Fatalf("critical control payload was written after normal payload: critical=%d normal=%d", criticalAt, normalAt)
	}
}

func TestControlBatchSenderCancelsQueuedRequestWithoutStoppingSender(t *testing.T) {
	conn := newBlockingRecordingControlConn()
	sender := newControlBatchSender(newTransportBinding(conn, stario.NewQueue()))
	defer sender.stop()
	defer conn.unblock()

	firstDone := make(chan error, 1)
	go func() {
		firstDone <- sender.submitContext(context.Background(), []byte("control-first"), 0, controlPriorityNormal)
	}()
	select {
	case <-conn.startCh:
	case <-time.After(time.Second):
		t.Fatal("first control write did not start")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
	defer cancel()
	started := time.Now()
	err := sender.submitContext(ctx, []byte("control-canceled"), 0, controlPriorityNormal)
	if !errors.Is(err, os.ErrDeadlineExceeded) {
		t.Fatalf("queued submit error = %v, want legacy deadline match", err)
	}
	if time.Since(started) > 250*time.Millisecond {
		t.Fatalf("queued cancellation took too long: %s", time.Since(started))
	}

	conn.unblock()
	if err := <-firstDone; err != nil {
		t.Fatalf("first submit failed: %v", err)
	}
	if err := sender.submitContext(context.Background(), []byte("control-after-cancel"), 0, controlPriorityNormal); err != nil {
		t.Fatalf("sender stopped after queued cancellation: %v", err)
	}
	if bytes.Contains(conn.bytes(), []byte("control-canceled")) {
		t.Fatal("canceled queued control payload was written")
	}
}

func TestControlBatchSenderCancelsWhileWaitingForSharedWriteLock(t *testing.T) {
	binding := newTransportBinding(&serializedWriteTestConn{}, stario.NewQueue())
	sender := newControlBatchSender(binding)
	defer sender.stop()

	lockHeld := make(chan struct{})
	releaseLock := make(chan struct{})
	lockDone := make(chan error, 1)
	go func() {
		lockDone <- binding.withConnWriteLock(func(net.Conn) error {
			close(lockHeld)
			<-releaseLock
			return nil
		})
	}()
	select {
	case <-lockHeld:
	case <-time.After(time.Second):
		t.Fatal("shared transport write lock was not acquired")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
	defer cancel()
	errCh := make(chan error, 1)
	go func() {
		errCh <- sender.submitContext(ctx, []byte("control-lock-wait"), 0, controlPriorityNormal)
	}()

	var err error
	select {
	case err = <-errCh:
	case <-time.After(250 * time.Millisecond):
		close(releaseLock)
		<-lockDone
		t.Fatal("control submit ignored context while waiting for the shared write lock")
	}
	stage, ok := TransportSendErrorStage(err)
	if !ok || stage != TransportSendStageQueue || !errors.Is(err, os.ErrDeadlineExceeded) {
		close(releaseLock)
		<-lockDone
		t.Fatalf("control lock-wait error=%v stage=%q,%v; want queue deadline", err, stage, ok)
	}
	if sender.errSnapshot() != nil {
		close(releaseLock)
		<-lockDone
		t.Fatalf("caller cancellation stopped a healthy control sender: %v", sender.errSnapshot())
	}

	close(releaseLock)
	if err := <-lockDone; err != nil {
		t.Fatalf("shared transport lock holder returned error: %v", err)
	}
	if err := sender.submitContext(context.Background(), []byte("control-after-lock-cancel"), 0, controlPriorityNormal); err != nil {
		t.Fatalf("sender failed after lock-wait cancellation: %v", err)
	}
}

func TestControlBatchSenderIsNotStarvedByContinuousSharedWriters(t *testing.T) {
	binding := newTransportBinding(&serializedWriteTestConn{}, stario.NewQueue())
	sender := newControlBatchSender(binding)
	defer sender.stop()

	if err := binding.lockConnWriteContext(context.Background()); err != nil {
		t.Fatalf("initial shared transport lock failed: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
	defer cancel()
	controlDone := make(chan error, 1)
	go func() {
		controlDone <- sender.submitContext(ctx, []byte("control-fairness"), 0, controlPriorityCritical)
	}()

	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		if sender.flushMu.TryLock() {
			sender.flushMu.Unlock()
			time.Sleep(time.Millisecond)
			continue
		}
		break
	}
	if sender.flushMu.TryLock() {
		sender.flushMu.Unlock()
		binding.unlockConnWrite()
		t.Fatal("control request did not enter the shared write wait")
	}

	var writers sync.WaitGroup
	for i := 0; i < 64; i++ {
		writers.Add(1)
		go func() {
			defer writers.Done()
			_ = binding.withConnWriteLock(func(net.Conn) error {
				time.Sleep(3 * time.Millisecond)
				return nil
			})
		}()
	}
	time.Sleep(10 * time.Millisecond)
	binding.unlockConnWrite()

	err := <-controlDone
	writers.Wait()
	if err != nil {
		t.Fatalf("control request was starved by shared writers: %v", err)
	}
}

func TestControlBatchStartedWriteCancellationReturnsBeforePhysicalWrite(t *testing.T) {
	conn := newBlockingRecordingControlConn()
	sender := newControlBatchSender(newTransportBinding(conn, stario.NewQueue()))
	defer sender.stop()
	defer conn.unblock()

	ctx, cancel := context.WithCancel(context.Background())
	payload := []byte("control-started-owned")
	done := make(chan error, 1)
	go func() {
		done <- sender.submitContext(ctx, payload, 0, controlPriorityCritical)
	}()
	select {
	case <-conn.startCh:
	case <-time.After(time.Second):
		t.Fatal("control write did not start")
	}

	cancel()
	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("started control cancellation error=%v, want context canceled", err)
		}
	case <-time.After(150 * time.Millisecond):
		t.Fatal("started control write kept the caller blocked after cancellation")
	}
	for i := range payload {
		payload[i] = 'x'
	}
	if sender.queued.Load() == 0 || sender.queuedSize.Load() == 0 {
		t.Fatal("started write released queue ownership before the physical write completed")
	}

	conn.unblock()
	deadline := time.Now().Add(time.Second)
	for (sender.queued.Load() != 0 || sender.queuedSize.Load() != 0) && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if got := sender.queued.Load(); got != 0 {
		t.Fatalf("queued requests after physical write = %d, want 0", got)
	}
	if got := sender.queuedSize.Load(); got != 0 {
		t.Fatalf("queued bytes after physical write = %d, want 0", got)
	}
	if !bytes.Contains(conn.bytes(), []byte("control-started-owned")) {
		t.Fatal("sender did not retain the started payload after caller cancellation")
	}
	if err := sender.submit([]byte("control-after-started-cancel"), 0); err != nil {
		t.Fatalf("sender failed after a started caller cancellation: %v", err)
	}
}

func TestControlBatchSenderStopCancelsDirectWaitForSharedWriteLock(t *testing.T) {
	binding := newTransportBinding(&serializedWriteTestConn{}, stario.NewQueue())
	sender := newControlBatchSender(binding)
	lockHeld := make(chan struct{})
	releaseLock := make(chan struct{})
	var releaseOnce sync.Once
	release := func() { releaseOnce.Do(func() { close(releaseLock) }) }
	defer sender.stop()
	defer release()

	lockDone := make(chan error, 1)
	go func() {
		lockDone <- binding.withConnWriteLock(func(net.Conn) error {
			close(lockHeld)
			<-releaseLock
			return nil
		})
	}()
	select {
	case <-lockHeld:
	case <-time.After(time.Second):
		t.Fatal("shared transport write lock was not acquired")
	}

	submitDone := make(chan error, 1)
	go func() {
		submitDone <- sender.submitContext(context.Background(), []byte("control-stop-lock-wait"), 0, controlPriorityNormal)
	}()
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		if sender.flushMu.TryLock() {
			sender.flushMu.Unlock()
			time.Sleep(time.Millisecond)
			continue
		}
		break
	}
	if sender.flushMu.TryLock() {
		sender.flushMu.Unlock()
		t.Fatal("direct control request did not enter the shared-lock wait")
	}

	stopDone := make(chan struct{})
	go func() {
		sender.stop()
		close(stopDone)
	}()
	select {
	case <-stopDone:
	case <-time.After(250 * time.Millisecond):
		release()
		<-lockDone
		<-submitDone
		<-stopDone
		t.Fatal("sender stop did not cancel a direct request waiting for the shared write lock")
	}
	select {
	case err := <-submitDone:
		if !errors.Is(err, errTransportDetached) {
			t.Fatalf("direct submit after sender stop error=%v, want transport detached", err)
		}
	case <-time.After(time.Second):
		t.Fatal("direct submit did not return after sender stop")
	}

	release()
	if err := <-lockDone; err != nil {
		t.Fatalf("shared transport lock holder returned error: %v", err)
	}
}

func TestControlBatchSenderRejectsOversizedPayloadBeforeWrite(t *testing.T) {
	conn := &serializedWriteTestConn{}
	sender := newControlBatchSender(newTransportBinding(conn, stario.NewQueue()))
	defer sender.stop()

	err := sender.submitContext(
		context.Background(),
		make([]byte, controlBatchMaxPayloadBytes+1),
		0,
		controlPriorityNormal,
	)
	stage, ok := TransportSendErrorStage(err)
	if !ok || stage != TransportSendStageQueue {
		t.Fatalf("oversized submit error=%v stage=%q,%v; want queue-stage rejection", err, stage, ok)
	}
	if err == nil || !strings.Contains(err.Error(), "stream or bulk") {
		t.Fatalf("oversized submit error=%v; want stream or bulk guidance", err)
	}
	if got := atomic.LoadInt32(&conn.writeCount); got != 0 {
		t.Fatalf("oversized control payload performed %d physical writes, want 0", got)
	}
	if sender.errSnapshot() != nil {
		t.Fatalf("oversized caller payload stopped a healthy sender: %v", sender.errSnapshot())
	}
	if err := sender.submit([]byte("control-after-oversized"), 0); err != nil {
		t.Fatalf("sender failed after oversized payload rejection: %v", err)
	}
}

func TestControlBatchSenderBoundsCriticalBurstWhenNormalWaits(t *testing.T) {
	conn := newBlockingRecordingControlConn()
	conn.unblock()
	sender := newControlBatchSender(newTransportBinding(conn, stario.NewQueue()))
	sender.flushMu.Lock()
	flushLocked := true
	defer func() {
		if flushLocked {
			sender.flushMu.Unlock()
		}
		sender.stop()
	}()

	firstCriticalDone := make(chan error, 1)
	go func() {
		firstCriticalDone <- sender.submitContext(context.Background(), []byte("control-critical-first"), 0, controlPriorityCritical)
	}()
	deadline := time.Now().Add(time.Second)
	for sender.queued.Load() != 1 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if sender.queued.Load() != 1 {
		t.Fatal("first critical request did not enter the queued path")
	}
	// Let run consume this request as a one-item batch before adding the rest.
	time.Sleep(20 * time.Millisecond)

	normalDone := make(chan error, 1)
	go func() {
		normalDone <- sender.submitContext(context.Background(), []byte("control-normal-waiting"), 0, controlPriorityNormal)
	}()
	criticalDone := make([]chan error, controlBatchMaxCriticalBurst)
	for i := range criticalDone {
		criticalDone[i] = make(chan error, 1)
		payload := []byte(fmt.Sprintf("control-critical-%02d", i))
		go func(done chan<- error, payload []byte) {
			done <- sender.submitContext(context.Background(), payload, 0, controlPriorityCritical)
		}(criticalDone[i], payload)
	}
	deadline = time.Now().Add(time.Second)
	for sender.queued.Load() != int64(controlBatchMaxCriticalBurst+2) && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if got, want := sender.queued.Load(), int64(controlBatchMaxCriticalBurst+2); got != want {
		t.Fatalf("queued requests=%d, want %d before releasing the sender", got, want)
	}

	sender.flushMu.Unlock()
	flushLocked = false
	for name, done := range map[string]<-chan error{
		"first critical": firstCriticalDone,
		"normal":         normalDone,
	} {
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("%s submit failed: %v", name, err)
			}
		case <-time.After(2 * time.Second):
			t.Fatalf("%s submit did not finish", name)
		}
	}
	for i, done := range criticalDone {
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("critical submit %d failed: %v", i, err)
			}
		case <-time.After(2 * time.Second):
			t.Fatalf("critical submit %d did not finish", i)
		}
	}

	written := conn.bytes()
	normalAt := bytes.Index(written, []byte("control-normal-waiting"))
	if normalAt < 0 {
		t.Fatal("normal control payload was not written")
	}
	criticalBeforeNormal := bytes.Count(written[:normalAt], []byte("control-critical-"))
	if criticalBeforeNormal > controlBatchMaxCriticalBurst {
		t.Fatalf("normal request waited behind %d critical messages, max burst is %d", criticalBeforeNormal, controlBatchMaxCriticalBurst)
	}
}

func TestControlBatchByteBudgetDoesNotAppendPastSoftLimit(t *testing.T) {
	limit := 4 * 1024 * 1024
	if !controlBatchCanAppend(1024*1024, 2*1024*1024, limit) {
		t.Fatal("batch should accept request below the cumulative byte limit")
	}
	if controlBatchCanAppend(3*1024*1024, 2*1024*1024, limit) {
		t.Fatal("batch accepted request past the cumulative byte limit")
	}
	if !controlBatchCanAppend(0, 8*1024*1024, limit) {
		t.Fatal("one oversized request must remain sendable as its own batch")
	}
}

func TestControlBatchSubmitReturnsWhenSenderStopsAfterQueue(t *testing.T) {
	sender := &controlBatchSender{
		normalCh:   make(chan controlBatchRequest, 1),
		criticalCh: make(chan controlBatchRequest, 1),
		stopCh:     make(chan struct{}),
		doneCh:     make(chan struct{}),
		budgetWake: make(chan struct{}, 1),
	}
	sender.flushMu.Lock()
	defer sender.flushMu.Unlock()

	errCh := make(chan error, 1)
	go func() {
		errCh <- sender.submitContext(context.Background(), []byte("queued-before-stop"), 0, controlPriorityNormal)
	}()
	deadline := time.Now().Add(time.Second)
	for len(sender.normalCh) == 0 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if len(sender.normalCh) == 0 {
		t.Fatal("control request did not enter the queue")
	}
	sender.markFailed(newTransportSendError(TransportSendStageTransport, errTransportDetached))

	select {
	case err := <-errCh:
		if err == nil {
			t.Fatal("queued submit returned nil after sender stop")
		}
		stage, ok := TransportSendErrorStage(err)
		if !ok || stage != TransportSendStageTransport {
			t.Fatalf("queued submit error stage=%q, %v; want %q, true", stage, ok, TransportSendStageTransport)
		}
	case <-time.After(time.Second):
		t.Fatal("queued submit remained blocked after sender stop")
	}
}

func TestControlBatchSenderStopAdmissionReleasesAllQueueBudget(t *testing.T) {
	for iteration := 0; iteration < 100; iteration++ {
		sender := newControlBatchSender(newTransportBinding(&serializedWriteTestConn{}, stario.NewQueue()))
		start := make(chan struct{})
		results := make(chan error, 32)
		for i := 0; i < cap(results); i++ {
			go func() {
				<-start
				results <- sender.submit([]byte("stop-admission-race"), 0)
			}()
		}
		close(start)
		sender.markFailed(newTransportSendError(TransportSendStageTransport, errTransportDetached))
		sender.stop()
		for i := 0; i < cap(results); i++ {
			<-results
		}
		if got := sender.queued.Load(); got != 0 {
			t.Fatalf("iteration %d queued requests=%d after stop, want 0", iteration, got)
		}
		if got := sender.queuedSize.Load(); got != 0 {
			t.Fatalf("iteration %d queued bytes=%d after stop, want 0", iteration, got)
		}
	}
}

func TestTransportSendErrorPreservesStageAndCause(t *testing.T) {
	err := newTransportSendError(TransportSendStageQueue, context.DeadlineExceeded)
	if !errors.Is(err, os.ErrDeadlineExceeded) {
		t.Fatalf("transport send error does not preserve legacy deadline matching: %v", err)
	}
	stage, ok := TransportSendErrorStage(err)
	if !ok || stage != TransportSendStageQueue {
		t.Fatalf("transport send stage = %q, %v; want %q, true", stage, ok, TransportSendStageQueue)
	}
}

func TestPublicContextSendErrorRestoresNormalizedDeadlineSentinel(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
	defer cancel()
	<-ctx.Done()

	err := newTransportSendError(TransportSendStageQueue, ctx.Err())
	if got := publicContextSendError(ctx, err); got != os.ErrDeadlineExceeded {
		t.Fatalf("public context send error=%#v, want exact os.ErrDeadlineExceeded", got)
	}
}

type blockingPacketWriteConn struct {
	startCh    chan struct{}
	unblockCh  chan struct{}
	writeCount atomic.Int32
}

func newBlockingPacketWriteConn() *blockingPacketWriteConn {
	return &blockingPacketWriteConn{
		startCh:   make(chan struct{}),
		unblockCh: make(chan struct{}),
	}
}

func (c *blockingPacketWriteConn) Read([]byte) (int, error) { return 0, net.ErrClosed }
func (c *blockingPacketWriteConn) Close() error             { return nil }
func (c *blockingPacketWriteConn) LocalAddr() net.Addr {
	return &net.UDPAddr{IP: net.IPv4zero, Port: 1}
}
func (c *blockingPacketWriteConn) RemoteAddr() net.Addr {
	return &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 2}
}
func (c *blockingPacketWriteConn) SetDeadline(time.Time) error      { return nil }
func (c *blockingPacketWriteConn) SetReadDeadline(time.Time) error  { return nil }
func (c *blockingPacketWriteConn) SetWriteDeadline(time.Time) error { return nil }

func (c *blockingPacketWriteConn) Write(p []byte) (int, error) {
	if c.writeCount.Add(1) == 1 {
		close(c.startCh)
		<-c.unblockCh
	}
	return len(p), nil
}

type vectoredShortWriteConn struct {
	steps  []int64
	idx    int
	buf    bytes.Buffer
	writes int
	writev int
}

func (c *vectoredShortWriteConn) Read([]byte) (int, error) { return 0, io.EOF }
func (c *vectoredShortWriteConn) Write(p []byte) (int, error) {
	c.writes++
	return c.buf.Write(p)
}
func (c *vectoredShortWriteConn) Close() error                     { return nil }
func (c *vectoredShortWriteConn) LocalAddr() net.Addr              { return nil }
func (c *vectoredShortWriteConn) RemoteAddr() net.Addr             { return nil }
func (c *vectoredShortWriteConn) SetDeadline(time.Time) error      { return nil }
func (c *vectoredShortWriteConn) SetReadDeadline(time.Time) error  { return nil }
func (c *vectoredShortWriteConn) SetWriteDeadline(time.Time) error { return nil }

func (c *vectoredShortWriteConn) WriteBuffers(bufs *net.Buffers) (int64, error) {
	c.writev++
	if c.idx >= len(c.steps) {
		return 0, io.ErrNoProgress
	}
	remaining := c.steps[c.idx]
	c.idx++
	written := int64(0)
	for len(*bufs) > 0 && remaining > 0 {
		part := (*bufs)[0]
		if len(part) == 0 {
			(*bufs)[0] = nil
			*bufs = (*bufs)[1:]
			continue
		}
		n := int64(len(part))
		if n > remaining {
			n = remaining
		}
		_, _ = c.buf.Write(part[:n])
		written += n
		remaining -= n
		if n == int64(len(part)) {
			(*bufs)[0] = nil
			*bufs = (*bufs)[1:]
			continue
		}
		(*bufs)[0] = part[n:]
		break
	}
	return written, nil
}

func (c *vectoredShortWriteConn) writeBuffers(bufs *net.Buffers) (int64, error) {
	return c.WriteBuffers(bufs)
}

type unwrapVectoredConn struct {
	inner net.Conn
}

func (c *unwrapVectoredConn) Read(p []byte) (int, error)         { return c.inner.Read(p) }
func (c *unwrapVectoredConn) Write(p []byte) (int, error)        { return c.inner.Write(p) }
func (c *unwrapVectoredConn) Close() error                       { return c.inner.Close() }
func (c *unwrapVectoredConn) LocalAddr() net.Addr                { return c.inner.LocalAddr() }
func (c *unwrapVectoredConn) RemoteAddr() net.Addr               { return c.inner.RemoteAddr() }
func (c *unwrapVectoredConn) SetDeadline(t time.Time) error      { return c.inner.SetDeadline(t) }
func (c *unwrapVectoredConn) SetReadDeadline(t time.Time) error  { return c.inner.SetReadDeadline(t) }
func (c *unwrapVectoredConn) SetWriteDeadline(t time.Time) error { return c.inner.SetWriteDeadline(t) }
func (c *unwrapVectoredConn) UnwrapConn() net.Conn               { return c.inner }

func TestWriteNetBuffersFullUnlockedFallsBackToDirectWritesAfterFirstPartialVectoredWrite(t *testing.T) {
	conn := &vectoredShortWriteConn{steps: []int64{3}}
	header := []byte("head")
	payload := []byte("payload")
	if err := writeNetBuffersFullUnlocked(conn, net.Buffers{header, payload}); err != nil {
		t.Fatalf("writeNetBuffersFullUnlocked failed: %v", err)
	}
	if got, want := conn.writev, 1; got != want {
		t.Fatalf("vectored write calls = %d, want %d", got, want)
	}
	if got, want := conn.writes, 2; got != want {
		t.Fatalf("fallback direct writes = %d, want %d", got, want)
	}
	if got, want := conn.buf.String(), "headpayload"; got != want {
		t.Fatalf("written bytes = %q, want %q", got, want)
	}
}

func TestWriteNetBuffersFullUnlockedReturnsNoProgressWhenVectoredWriteDoesNotAdvance(t *testing.T) {
	conn := &vectoredShortWriteConn{steps: []int64{0}}
	err := writeNetBuffersFullUnlocked(conn, net.Buffers{[]byte("head"), []byte("payload")})
	if !errors.Is(err, io.ErrNoProgress) {
		t.Fatalf("writeNetBuffersFullUnlocked error = %v, want %v", err, io.ErrNoProgress)
	}
}

func TestWriteNetBuffersFullUnlockedUsesUnwrappedVectoredConn(t *testing.T) {
	inner := &vectoredShortWriteConn{steps: []int64{100}}
	conn := &unwrapVectoredConn{inner: inner}
	if err := writeNetBuffersFullUnlocked(conn, net.Buffers{[]byte("head"), []byte("payload")}); err != nil {
		t.Fatalf("writeNetBuffersFullUnlocked failed: %v", err)
	}
	if got, want := inner.writev, 1; got != want {
		t.Fatalf("unwrapped vectored write calls = %d, want %d", got, want)
	}
	if got := inner.writes; got != 0 {
		t.Fatalf("unexpected fallback direct writes = %d, want 0", got)
	}
}

func TestBulkBatchSenderSkipsQueuedCanceledRequest(t *testing.T) {
	conn := newBlockingPacketWriteConn()
	binding := newTransportBinding(conn, stario.NewQueue())
	sender := newTestBulkBatchSender(binding)
	defer sender.stop()

	firstErrCh := make(chan error, 1)
	go func() {
		firstErrCh <- sender.submitData(context.Background(), 1, 1, bulkFastPathVersionV1, []byte("first"))
	}()

	select {
	case <-conn.startCh:
	case <-time.After(time.Second):
		t.Fatal("first shared bulk write did not start")
	}

	ctx, cancel := context.WithCancel(context.Background())
	secondErrCh := make(chan error, 1)
	go func() {
		secondErrCh <- sender.submitData(ctx, 1, 2, bulkFastPathVersionV1, []byte("second"))
	}()
	time.Sleep(20 * time.Millisecond)
	cancel()

	select {
	case err := <-secondErrCh:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("second shared bulk submit error = %v, want %v", err, context.Canceled)
		}
	case <-time.After(time.Second):
		t.Fatal("second shared bulk submit did not return after cancel")
	}

	close(conn.unblockCh)

	select {
	case err := <-firstErrCh:
		if err != nil {
			t.Fatalf("first shared bulk submit failed: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("first shared bulk submit did not finish")
	}

	time.Sleep(50 * time.Millisecond)
	if got, want := conn.writeCount.Load(), int32(1); got != want {
		t.Fatalf("shared bulk write count = %d, want %d", got, want)
	}
}

func TestBulkBatchSenderDoesNotDirectSubmitShareableV2Data(t *testing.T) {
	sender := &bulkBatchSender{}
	req := bulkBatchRequest{
		frames: []bulkFastFrame{{
			Type:    bulkFastPayloadTypeData,
			DataID:  1,
			Seq:     1,
			Payload: make([]byte, 256*1024),
		}},
		fastPathVersion: bulkFastPathVersionV2,
	}
	if sender.shouldDirectSubmit(req) {
		t.Fatal("shareable v2 shared bulk data should queue for super-batch instead of direct submit")
	}
}

func TestBulkBatchSenderDirectSubmitsUnbatchableRequest(t *testing.T) {
	sender := &bulkBatchSender{}
	req := bulkBatchRequest{
		frames: []bulkFastFrame{{
			Type:    bulkFastPayloadTypeRelease,
			DataID:  1,
			Seq:     0,
			Payload: []byte("rel"),
		}},
		fastPathVersion: bulkFastPathVersionV2,
	}
	if !sender.shouldDirectSubmit(req) {
		t.Fatal("unbatchable shared bulk control request should still direct submit")
	}
}

func TestBulkBatchSenderReturnsFlushResultAfterStartedContextCancel(t *testing.T) {
	conn := newBlockingPacketWriteConn()
	binding := newTransportBinding(conn, stario.NewQueue())
	sender := newTestBulkBatchSender(binding)
	defer sender.stop()

	ctx, cancel := context.WithCancel(context.Background())
	errCh := make(chan error, 1)
	go func() {
		errCh <- sender.submitData(ctx, 1, 1, bulkFastPathVersionV1, []byte("payload"))
	}()

	select {
	case <-conn.startCh:
	case <-time.After(time.Second):
		t.Fatal("shared bulk write did not start")
	}

	cancel()

	select {
	case err := <-errCh:
		t.Fatalf("sender.submit returned before flush completed: %v", err)
	case <-time.After(50 * time.Millisecond):
	}

	close(conn.unblockCh)

	select {
	case err := <-errCh:
		if err != nil {
			t.Fatalf("sender.submit failed after started flush: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("sender.submit did not return after started flush completed")
	}
}

func TestTransportBindingStopBackgroundWorkersStopsSharedSender(t *testing.T) {
	binding := newTransportBinding(newBlockingPacketWriteConn(), stario.NewQueue())
	sender := binding.bulkBatchSenderSnapshotWithCodec(bulkBatchCodec{
		encodeSingle: func(frame bulkFastFrame) ([]byte, func(), error) {
			return append([]byte(nil), frame.Payload...), nil, nil
		},
		encodeBatch: func(frames []bulkFastFrame) ([]byte, func(), error) {
			payload, err := encodeBulkFastBatchPlain(frames)
			return payload, nil, err
		},
	}, nil)
	binding.stopBackgroundWorkers()

	err := sender.submitData(context.Background(), 1, 1, bulkFastPathVersionV1, []byte("payload"))
	if !errors.Is(err, errTransportDetached) {
		t.Fatalf("sender.submit after stop = %v, want %v", err, errTransportDetached)
	}
}

func TestTransportBindingStopBackgroundWorkersStopsControlSender(t *testing.T) {
	binding := newTransportBinding(&serializedWriteTestConn{}, stario.NewQueue())
	sender := binding.controlBatchSenderSnapshot()
	binding.stopBackgroundWorkers()

	err := sender.submit([]byte("payload"), 0)
	if !errors.Is(err, errTransportDetached) {
		t.Fatalf("sender.submit after stop = %v, want %v", err, errTransportDetached)
	}
}