diff --git a/internal/grid/connection.go b/internal/grid/connection.go index a9eca99df..f761b9389 100644 --- a/internal/grid/connection.go +++ b/internal/grid/connection.go @@ -287,6 +287,7 @@ func (c *Connection) newMuxClient(ctx context.Context) (*muxClient, error) { if dl, ok := ctx.Deadline(); ok { client.deadline = getDeadline(time.Until(dl)) if client.deadline == 0 { + client.cancelFn(context.DeadlineExceeded) return nil, context.DeadlineExceeded } } @@ -333,6 +334,7 @@ func (c *Connection) Request(ctx context.Context, h HandlerID, req []byte) ([]by _, ok := c.outgoing.Load(client.MuxID) fmt.Println(client.MuxID, c.String(), "Connection.Request: DELETING MUX. Exists:", ok) } + client.cancelFn(context.Canceled) c.outgoing.Delete(client.MuxID) }() return client.traceRoundtrip(ctx, c.trace, h, req) diff --git a/internal/grid/stream.go b/internal/grid/stream.go index abebd2229..d65313d77 100644 --- a/internal/grid/stream.go +++ b/internal/grid/stream.go @@ -63,10 +63,11 @@ func (s *Stream) Send(b []byte) error { func (s *Stream) Results(next func(b []byte) error) (err error) { done := false defer func() { + if s.cancel != nil { + s.cancel(err) + } + if !done { - if s.cancel != nil { - s.cancel(err) - } // Drain channel. for range s.responses { }