From 32df3f7c5ba29b1d23426c37ae75b09c961cd355 Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Mon, 27 Jul 2026 14:11:13 -0500 Subject: [PATCH 1/2] fix(gravity): bound tunnel sends and isolate stalled endpoints --- gravity/grpc_client.go | 163 ++++++++++++++++++++++--------- gravity/tunnel_dataplane_test.go | 118 ++++++++++++++++++++++ 2 files changed, 236 insertions(+), 45 deletions(-) diff --git a/gravity/grpc_client.go b/gravity/grpc_client.go index 5fd77a5..dc95074 100644 --- a/gravity/grpc_client.go +++ b/gravity/grpc_client.go @@ -63,6 +63,9 @@ func init() { // Error variables for consistency with old implementation var ErrConnectionClosed = errors.New("gravity connection closed") +var ErrTunnelSendTimeout = errors.New("gravity tunnel send timed out") + +const defaultTunnelSendTimeout = 5 * time.Second const ( // debugTunnelPackets gates verbose per-packet tunnel logging. @@ -296,6 +299,7 @@ type GravityClient struct { onEndpointReady func(endpointIndex int) onEndpointTunnelReady func(endpointIndex int) reconnectEndpointHook func(endpointIndex int, reason string) // test hook; nil in production + tunnelSendTimeout time.Duration // State management connected bool @@ -429,12 +433,18 @@ func sessionHelloErrorMessage(err error) string { // StreamInfo tracks individual stream health and load type StreamInfo struct { stream pb.GravitySessionService_StreamSessionPacketsClient - connIndex int // Connection index this stream belongs to - streamID string // Unique stream identifier - isHealthy bool // Stream health status - loadCount int64 // Current load (packets in flight) - lastUsed time.Time // Last time this stream was used - sendMu sync.Mutex // Serializes Send calls on this stream + connIndex int // Connection index this stream belongs to + streamID string // Unique stream identifier + isHealthy bool // Stream health status + loadCount int64 // Current load (packets in flight) + lastUsed time.Time // Last time this stream was used + sendOnce sync.Once + sendQueue chan tunnelSendRequest +} + +type tunnelSendRequest struct { + packet *pb.TunnelPacket + result chan error } const endpointDiscoveryRefreshFailureThreshold int32 = 10 @@ -569,6 +579,7 @@ func New(config GravityConfig) (*GravityClient, error) { maxReconnectAttempts: config.MaxReconnectAttempts, reconnectAttemptTimeout: config.ReconnectAttemptTimeout, reconnectionFailedCallback: config.ReconnectionFailedCallback, + tunnelSendTimeout: defaultTunnelSendTimeout, onEndpointReady: config.OnEndpointReady, onEndpointTunnelReady: config.OnEndpointTunnelReady, tracer: otel.Tracer("@agentuity/gravity/client"), @@ -1368,7 +1379,7 @@ type streamSnapshot struct { connIndex int isHealthy bool stream pb.GravitySessionService_StreamSessionPacketsClient - sendMu *sync.Mutex // pointer to the stream's send mutex for serialization + info *StreamInfo } func (g *GravityClient) sendTunnelKeepalives() { @@ -1383,7 +1394,7 @@ func (g *GravityClient) sendTunnelKeepalives() { connIndex: si.connIndex, isHealthy: si.isHealthy, stream: si.stream, - sendMu: &si.sendMu, + info: si, } } } @@ -1427,7 +1438,7 @@ func (g *GravityClient) sendTunnelKeepalives() { healthyInEndpoint := make([]int, 0, streamsPerConn) for i := base; i < end; i++ { s := snapshots[i] - if s.isHealthy && s.stream != nil && s.sendMu != nil { + if s.isHealthy && s.stream != nil && s.info != nil { healthyInEndpoint = append(healthyInEndpoint, i) } } @@ -1439,18 +1450,11 @@ func (g *GravityClient) sendTunnelKeepalives() { targetIdx := healthyInEndpoint[rotation%len(healthyInEndpoint)] s := snapshots[targetIdx] - // Serialize with the data path's sendMu. Use a goroutine with - // timeout so a wedged Send doesn't block the monitor loop. + // Serialize with the data path's sendMu. A timed-out send quarantines + // and reconnects only this endpoint. go func(idx int, snap streamSnapshot) { - errCh := make(chan error, 1) - go func() { - snap.sendMu.Lock() - err := snap.stream.Send(&pb.TunnelPacket{Data: probe}) - snap.sendMu.Unlock() - errCh <- err - }() - select { - case err := <-errCh: + err := g.sendTunnelStreamBounded(snap.info, &pb.TunnelPacket{Data: probe}, "keepalive") + if err == nil { // Update LastProbeSentUs on every successful probe so the // liveness monitor can distinguish "probed but no echo" // from "never probed." Unlike LastSendUs (updated by @@ -1458,15 +1462,13 @@ func (g *GravityClient) sendTunnelKeepalives() { // set by keepalive probes — preventing data-path sends // from masking a zombie tunnel (one-way: sending works, // receiving doesn't). - if err == nil { - g.streamManager.metricsMu.Lock() - if m := g.streamManager.streamMetrics[snap.streamID]; m != nil { - m.LastProbeSentUs = time.Now().UnixMicro() - } - g.streamManager.metricsMu.Unlock() + g.streamManager.metricsMu.Lock() + if m := g.streamManager.streamMetrics[snap.streamID]; m != nil { + m.LastProbeSentUs = time.Now().UnixMicro() } - case <-time.After(5 * time.Second): - g.logger.Warn("tunnel keepalive send blocked >5s on stream %d — possible zombie", idx) + g.streamManager.metricsMu.Unlock() + } else { + g.logger.Warn("tunnel keepalive send failed on stream %d: %v", idx, err) } }(targetIdx, s) } @@ -3909,22 +3911,22 @@ func (g *GravityClient) stop() error { g.closing = true g.mu.Unlock() - // Wait for in-flight checkpoint/restore handlers to finish so their - // response messages reach the server before we tear down streams. - g.handlerWg.Wait() - g.mu.Lock() connectionCancel := g.connectionCancel cancel := g.cancel g.connected = false g.mu.Unlock() + // Cancel stream contexts before waiting for handlers. A handler or packet + // sender may be blocked in gRPC I/O; cancellation is what makes that I/O + // return and guarantees bounded shutdown. if connectionCancel != nil { connectionCancel() } if cancel != nil { cancel() } + g.handlerWg.Wait() g.cleanup() return nil @@ -4124,6 +4126,11 @@ func (g *GravityClient) handleEndpointDisconnection(endpointIndex int, reason st } else { g.logger.Warn("all endpoints down, reconnecting endpoint %d independently", endpointIndex) } + if g.reconnectEndpointHook != nil { + defer g.endpointReconnecting[endpointIndex].Store(false) + g.reconnectEndpointHook(endpointIndex, reason) + return + } // Always use per-endpoint reconnection in multi-endpoint mode. // Each endpoint reconnects independently as its gravity server // becomes available. This is better than full reconnect which @@ -5867,6 +5874,79 @@ func (g *GravityClient) Resume(reason string) error { } // WritePacket sends a tunnel packet via gRPC tunnel stream using load balancing +func (g *GravityClient) sendTunnelStreamBounded(stream *StreamInfo, packet *pb.TunnelPacket, reason string) error { + if stream == nil || stream.stream == nil { + return fmt.Errorf("tunnel stream is unavailable") + } + timeout := g.tunnelSendTimeout + if timeout <= 0 { + timeout = defaultTunnelSendTimeout + } + // Send may outlive the caller when the deadline fires. Give the send + // goroutine immutable ownership so a TUN buffer can be reused immediately. + ownedPacket := &pb.TunnelPacket{ + Data: append([]byte(nil), packet.Data...), + StreamId: packet.StreamId, + EnqueuedAtUs: packet.EnqueuedAtUs, + } + stream.sendOnce.Do(func() { + stream.sendQueue = make(chan tunnelSendRequest) + go g.runTunnelSendWorker(stream) + }) + req := tunnelSendRequest{packet: ownedPacket, result: make(chan error, 1)} + + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case stream.sendQueue <- req: + case <-g.ctx.Done(): + return ErrConnectionClosed + case <-timer.C: + g.quarantineTunnelSend(stream, reason, timeout) + return fmt.Errorf("%w after %s", ErrTunnelSendTimeout, timeout) + } + select { + case err := <-req.result: + return err + case <-g.ctx.Done(): + return ErrConnectionClosed + case <-timer.C: + g.quarantineTunnelSend(stream, reason, timeout) + return fmt.Errorf("%w after %s", ErrTunnelSendTimeout, timeout) + } +} + +func (g *GravityClient) runTunnelSendWorker(stream *StreamInfo) { + for { + select { + case <-g.ctx.Done(): + return + case <-stream.stream.Context().Done(): + return + case req := <-stream.sendQueue: + err := stream.stream.Send(req.packet) + req.result <- err + } + } +} + +func (g *GravityClient) quarantineTunnelSend(stream *StreamInfo, reason string, timeout time.Duration) { + // gRPC stream Send has no per-call context. Quarantine first so no new + // packets select this stream, then cancel/reconnect its endpoint. The + // stream context cancellation unblocks the outstanding Send worker. + g.streamManager.tunnelMu.Lock() + stream.isHealthy = false + g.streamManager.tunnelMu.Unlock() + g.endpointsMu.RLock() + if stream.connIndex >= 0 && stream.connIndex < len(g.endpoints) && g.endpoints[stream.connIndex] != nil { + g.endpoints[stream.connIndex].healthy.Store(false) + } + g.endpointsMu.RUnlock() + g.logger.Error("tunnel %s send timed out after %s on endpoint %d stream %s; quarantining endpoint", + reason, timeout, stream.connIndex, stream.streamID) + go g.handleEndpointDisconnection(stream.connIndex, "tunnel_"+reason+"_send_timeout") +} + func (g *GravityClient) WritePacket(payload []byte) error { if g.tracePackets { g.tracePacketLogger.Debug("writePacket called with %d bytes", len(payload)) @@ -5929,9 +6009,7 @@ func (g *GravityClient) WritePacket(payload []byte) error { } sendStart := time.Now() - stream.sendMu.Lock() - err = stream.stream.Send(tunnelPacket) - stream.sendMu.Unlock() + err = g.sendTunnelStreamBounded(stream, tunnelPacket, "packet") sendLatency := time.Since(sendStart) if err != nil { // Mark stream unhealthy and record per-stream error metrics, @@ -6198,21 +6276,16 @@ func (g *GravityClient) sendTunnelPacket(data []byte) error { StreamId: streamInfo.streamID, } - // Send packet with retry logic and circuit breaker + // Validate the endpoint's circuit breaker before sending. Tunnel sends are + // deliberately not retried here: retransmission belongs to the tunneled + // transport, and retrying a blocked gRPC Send extends the dataplane stall. connectionIndex := streamInfo.connIndex if connectionIndex < 0 || connectionIndex >= len(g.circuitBreakers) { return fmt.Errorf("circuit breaker index %d out of range (len=%d)", connectionIndex, len(g.circuitBreakers)) } - circuitBreaker := g.circuitBreakers[connectionIndex] sendStart := time.Now() - - err = RetryWithCircuitBreaker(context.WithoutCancel(g.ctx), g.retryConfig, circuitBreaker, func() error { - streamInfo.sendMu.Lock() - sendErr := streamInfo.stream.Send(packet) - streamInfo.sendMu.Unlock() - return sendErr - }) + err = g.sendTunnelStreamBounded(streamInfo, packet, "packet") sendLatency := time.Since(sendStart) @@ -6230,7 +6303,7 @@ func (g *GravityClient) sendTunnelPacket(data []byte) error { g.streamManager.metricsMu.Unlock() g.outboundErrors.Add(1) - return fmt.Errorf("failed to send packet after retries: %w", err) + return fmt.Errorf("failed to send packet: %w", err) } // Record successful send metrics diff --git a/gravity/tunnel_dataplane_test.go b/gravity/tunnel_dataplane_test.go index 8910ab3..47ae502 100644 --- a/gravity/tunnel_dataplane_test.go +++ b/gravity/tunnel_dataplane_test.go @@ -59,6 +59,8 @@ type dataplaneMockTunnelStream struct { errCh chan error closeSendCount atomic.Int64 + blockSend <-chan struct{} + activeSends atomic.Int64 } var _ pb.GravitySessionService_StreamSessionPacketsClient = (*dataplaneMockTunnelStream)(nil) @@ -75,6 +77,15 @@ func newDataplaneMockTunnelStream(ctx context.Context) *dataplaneMockTunnelStrea } func (m *dataplaneMockTunnelStream) Send(p *pb.TunnelPacket) error { + m.activeSends.Add(1) + defer m.activeSends.Add(-1) + if m.blockSend != nil { + select { + case <-m.blockSend: + case <-m.ctx.Done(): + return status.Error(codes.Canceled, "context canceled") + } + } if v := m.sendErr.Load(); v != nil { if err, ok := v.(error); ok && err != nil { return err @@ -224,6 +235,113 @@ func TestWritePacket_SendsToStream(t *testing.T) { } } +func TestWritePacket_BlockedSendIsBoundedAndReconnectsOnlyEndpoint(t *testing.T) { + g, _ := newTunnelDataplaneTestClient(t, 2) + g.tunnelSendTimeout = 25 * time.Millisecond + g.multiEndpointMode.Store(true) + g.selector = NewEndpointSelector(time.Minute) + + blockedCtx, blockedCancel := context.WithCancel(g.ctx) + defer blockedCancel() + blocked := newDataplaneMockTunnelStream(blockedCtx) + blocked.blockSend = make(chan struct{}) + sibling := newDataplaneMockTunnelStream(g.ctx) + installStreams(g, + &StreamInfo{stream: blocked, connIndex: 0, streamID: "blocked", isHealthy: true}, + &StreamInfo{stream: sibling, connIndex: 1, streamID: "sibling", isHealthy: true}, + ) + g.streamManager.controlStreams[0] = &configurableMockStream{} + g.streamManager.controlStreams[1] = &configurableMockStream{} + g.streamManager.cancels = []context.CancelFunc{blockedCancel, func() {}} + reconnected := make(chan int, 1) + g.reconnectEndpointHook = func(idx int, _ string) { reconnected <- idx } + + // Bind this flow to endpoint zero so the first send exercises the block. + g.selector.bindings[ExtractFlowKey(makeIPv6Packet())] = &TunnelBinding{Endpoint: g.endpoints[0], LastUsed: time.Now()} + start := time.Now() + err := g.WritePacket(makeIPv6Packet()) + if !errors.Is(err, ErrTunnelSendTimeout) { + t.Fatalf("expected tunnel send timeout, got %v", err) + } + if elapsed := time.Since(start); elapsed > 250*time.Millisecond { + t.Fatalf("blocked send was not bounded: %s", elapsed) + } + select { + case idx := <-reconnected: + if idx != 0 { + t.Fatalf("reconnected endpoint %d, want 0", idx) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for endpoint quarantine") + } + + if err := g.WritePacket(reverseFlow(makeIPv6Packet())); err != nil { + t.Fatalf("healthy sibling endpoint stopped carrying traffic: %v", err) + } + if sibling.sentCount() != 1 { + t.Fatalf("healthy sibling received %d packets, want 1", sibling.sentCount()) + } + if !g.endpoints[1].IsHealthy() { + t.Fatal("healthy sibling endpoint was quarantined") + } + waitUntil(t, time.Second, func() bool { return blocked.activeSends.Load() == 0 }) +} + +func TestKeepalive_BlockedSendQuarantinesEndpoint(t *testing.T) { + g, _ := newTunnelDataplaneTestClient(t, 2) + g.tunnelSendTimeout = 20 * time.Millisecond + blockedCtx, blockedCancel := context.WithCancel(g.ctx) + defer blockedCancel() + blocked := newDataplaneMockTunnelStream(blockedCtx) + blocked.blockSend = make(chan struct{}) + sibling := newDataplaneMockTunnelStream(g.ctx) + installStreams(g, + &StreamInfo{stream: blocked, connIndex: 0, streamID: "blocked", isHealthy: true}, + &StreamInfo{stream: sibling, connIndex: 1, streamID: "sibling", isHealthy: true}, + ) + g.streamManager.controlStreams[0] = &configurableMockStream{} + g.streamManager.controlStreams[1] = &configurableMockStream{} + g.poolConfig.StreamsPerGravity = 1 + g.streamManager.cancels = []context.CancelFunc{blockedCancel, func() {}} + reconnected := make(chan int, 1) + g.reconnectEndpointHook = func(idx int, _ string) { reconnected <- idx } + + g.sendTunnelKeepalives() + select { + case idx := <-reconnected: + if idx != 0 { + t.Fatalf("reconnected endpoint %d, want 0", idx) + } + case <-time.After(time.Second): + t.Fatal("blocked keepalive did not reconnect its endpoint") + } + waitUntil(t, time.Second, func() bool { return sibling.sentCount() == 1 }) + if !g.endpoints[1].IsHealthy() { + t.Fatal("sibling endpoint was quarantined by another endpoint's keepalive") + } + waitUntil(t, time.Second, func() bool { return blocked.activeSends.Load() == 0 }) +} + +func TestClose_CancelsHandlersBeforeWaiting(t *testing.T) { + g, _ := newTunnelDataplaneTestClient(t, 1) + g.handlerWg.Add(1) + go func() { + defer g.handlerWg.Done() + <-g.ctx.Done() + }() + + done := make(chan error, 1) + go func() { done <- g.Close() }() + select { + case err := <-done: + if err != nil { + t.Fatalf("Close failed: %v", err) + } + case <-time.After(time.Second): + t.Fatal("Close blocked waiting for a handler whose I/O context was not cancelled") + } +} + func TestSendPacketCopiesCallerBuffer(t *testing.T) { t.Parallel() g, _ := newTunnelDataplaneTestClient(t, 1) From 70e3174cd587c8250d6752346882958bc194e5a5 Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Mon, 27 Jul 2026 14:23:42 -0500 Subject: [PATCH 2/2] fix(gravity): handle exited tunnel send workers --- gravity/grpc_client.go | 50 ++++++++++++++++++++++++++++---- gravity/tunnel_dataplane_test.go | 27 +++++++++++++++++ 2 files changed, 71 insertions(+), 6 deletions(-) diff --git a/gravity/grpc_client.go b/gravity/grpc_client.go index dc95074..3bd89bf 100644 --- a/gravity/grpc_client.go +++ b/gravity/grpc_client.go @@ -440,11 +440,13 @@ type StreamInfo struct { lastUsed time.Time // Last time this stream was used sendOnce sync.Once sendQueue chan tunnelSendRequest + sendDone chan struct{} } type tunnelSendRequest struct { - packet *pb.TunnelPacket - result chan error + packet *pb.TunnelPacket + result chan error + pooledBuffer []byte } const endpointDiscoveryRefreshFailureThreshold int32 = 10 @@ -1450,8 +1452,8 @@ func (g *GravityClient) sendTunnelKeepalives() { targetIdx := healthyInEndpoint[rotation%len(healthyInEndpoint)] s := snapshots[targetIdx] - // Serialize with the data path's sendMu. A timed-out send quarantines - // and reconnects only this endpoint. + // Serialize through the data path's per-stream sendQueue worker. A + // timed-out send quarantines and reconnects only this endpoint. go func(idx int, snap streamSnapshot) { err := g.sendTunnelStreamBounded(snap.info, &pb.TunnelPacket{Data: probe}, "keepalive") if err == nil { @@ -5878,29 +5880,57 @@ func (g *GravityClient) sendTunnelStreamBounded(stream *StreamInfo, packet *pb.T if stream == nil || stream.stream == nil { return fmt.Errorf("tunnel stream is unavailable") } + select { + case <-g.ctx.Done(): + return ErrConnectionClosed + case <-stream.stream.Context().Done(): + return ErrConnectionClosed + default: + } timeout := g.tunnelSendTimeout if timeout <= 0 { timeout = defaultTunnelSendTimeout } // Send may outlive the caller when the deadline fires. Give the send // goroutine immutable ownership so a TUN buffer can be reused immediately. + var ownedData []byte + var pooledBuffer []byte + if len(packet.Data) <= maxBufferSize { + pooledBuffer = g.bufferPool.Get().([]byte) + ownedData = pooledBuffer[:len(packet.Data)] + copy(ownedData, packet.Data) + } else { + ownedData = append([]byte(nil), packet.Data...) + } ownedPacket := &pb.TunnelPacket{ - Data: append([]byte(nil), packet.Data...), + Data: ownedData, StreamId: packet.StreamId, EnqueuedAtUs: packet.EnqueuedAtUs, } stream.sendOnce.Do(func() { stream.sendQueue = make(chan tunnelSendRequest) + stream.sendDone = make(chan struct{}) go g.runTunnelSendWorker(stream) }) - req := tunnelSendRequest{packet: ownedPacket, result: make(chan error, 1)} + req := tunnelSendRequest{packet: ownedPacket, result: make(chan error, 1), pooledBuffer: pooledBuffer} + handedOff := false + defer func() { + if !handedOff && req.pooledBuffer != nil { + g.bufferPool.Put(req.pooledBuffer) + } + }() timer := time.NewTimer(timeout) defer timer.Stop() select { case stream.sendQueue <- req: + handedOff = true case <-g.ctx.Done(): return ErrConnectionClosed + case <-stream.stream.Context().Done(): + return ErrConnectionClosed + case <-stream.sendDone: + return ErrConnectionClosed case <-timer.C: g.quarantineTunnelSend(stream, reason, timeout) return fmt.Errorf("%w after %s", ErrTunnelSendTimeout, timeout) @@ -5910,6 +5940,10 @@ func (g *GravityClient) sendTunnelStreamBounded(stream *StreamInfo, packet *pb.T return err case <-g.ctx.Done(): return ErrConnectionClosed + case <-stream.stream.Context().Done(): + return ErrConnectionClosed + case <-stream.sendDone: + return ErrConnectionClosed case <-timer.C: g.quarantineTunnelSend(stream, reason, timeout) return fmt.Errorf("%w after %s", ErrTunnelSendTimeout, timeout) @@ -5917,6 +5951,7 @@ func (g *GravityClient) sendTunnelStreamBounded(stream *StreamInfo, packet *pb.T } func (g *GravityClient) runTunnelSendWorker(stream *StreamInfo) { + defer close(stream.sendDone) for { select { case <-g.ctx.Done(): @@ -5925,6 +5960,9 @@ func (g *GravityClient) runTunnelSendWorker(stream *StreamInfo) { return case req := <-stream.sendQueue: err := stream.stream.Send(req.packet) + if req.pooledBuffer != nil { + g.bufferPool.Put(req.pooledBuffer) + } req.result <- err } } diff --git a/gravity/tunnel_dataplane_test.go b/gravity/tunnel_dataplane_test.go index 47ae502..159dae4 100644 --- a/gravity/tunnel_dataplane_test.go +++ b/gravity/tunnel_dataplane_test.go @@ -287,6 +287,33 @@ func TestWritePacket_BlockedSendIsBoundedAndReconnectsOnlyEndpoint(t *testing.T) waitUntil(t, time.Second, func() bool { return blocked.activeSends.Load() == 0 }) } +func TestSendTunnelStreamBounded_ExitedWorkerReturnsConnectionClosedImmediately(t *testing.T) { + g, _ := newTunnelDataplaneTestClient(t, 1) + g.tunnelSendTimeout = time.Second + streamCtx, cancel := context.WithCancel(g.ctx) + mock := newDataplaneMockTunnelStream(streamCtx) + stream := &StreamInfo{stream: mock, connIndex: 0, streamID: "closed", isHealthy: true} + + if err := g.sendTunnelStreamBounded(stream, &pb.TunnelPacket{Data: []byte("first")}, "packet"); err != nil { + t.Fatalf("initial send failed: %v", err) + } + cancel() + select { + case <-stream.sendDone: + case <-time.After(time.Second): + t.Fatal("send worker did not exit") + } + + start := time.Now() + err := g.sendTunnelStreamBounded(stream, &pb.TunnelPacket{Data: []byte("second")}, "packet") + if !errors.Is(err, ErrConnectionClosed) { + t.Fatalf("expected ErrConnectionClosed, got %v", err) + } + if elapsed := time.Since(start); elapsed > 100*time.Millisecond { + t.Fatalf("closed stream waited for send timeout: %s", elapsed) + } +} + func TestKeepalive_BlockedSendQuarantinesEndpoint(t *testing.T) { g, _ := newTunnelDataplaneTestClient(t, 2) g.tunnelSendTimeout = 20 * time.Millisecond