Potential Goroutine Leak in `streamableClientConn`

[cryo-zd]

Description

I have observed a potential issue with goroutine leaks in the streamableClientConn implementation when handling Server-Sent Events (SSE). Below is a description of the issue and the potential scenario that may lead to a goroutine leak.

go-sdk/mcp/streamable.go

Lines 751 to 770 in c037ba5

	func (s *streamableClientConn) handleSSE(resp *http.Response) {
	defer resp.Body.Close()
	done := make(chan struct{})
	go func() {
	defer close(done)
	for evt, err := range scanEvents(resp.Body) {
	if err != nil {
	// TODO: surface this error; possibly break the stream
	return
	}
	s.incoming <- evt.data
	}
	}()
	select {
	case <-s.done:
	case <-done:
	}
	}

Problem Scenario

1. SSE Event Handling:
   When the server sends SSE events, a goroutine is spawned to read events and write them to the s.incoming channel:
   

   go-sdk/mcp/streamable.go

   Lines 744 to 748 in 6c6243c

   	if resp.Header.Get("Content-Type") == "text/event-stream" {
   	go s.handleSSE(resp)
   	} else {
   	resp.Body.Close()
   	}

2. Read() as the Sole Consumer:
   The Read() method is the only consumer of s.incoming. If Read() exits (due to s.done being closed by Close()), there will be no remaining consumer for s.incoming.
   

   go-sdk/mcp/streamable.go

   Lines 665 to 674 in 6c6243c

   	func (s *streamableClientConn) Read(ctx context.Context) (jsonrpc.Message, error) {
   	select {
   	case <-ctx.Done():
   	return nil, ctx.Err()
   	case <-s.done:
   	return nil, io.EOF
   	case data := <-s.incoming:
   	return jsonrpc2.DecodeMessage(data)
   	}
   	}

3. Blocked Goroutine:
   If s.incoming becomes full (due to slow consumption or high-frequency events), the goroutine writing to s.incoming will block and will not be able to proceed.
4. Close() Function Call:
   After Close() is called, s.done is closed, which causes the goroutine running handleSSE to exit. However, the sub-goroutine writing to s.incoming remains blocked because the channel is full. Since resp.Body.Close() is deferred in handleSSE, it is executed when the handleSSE goroutine exits, which leads to the closure of the underlying stream.
5. Goroutine Leak:
   Although resp.Body.Close() will cause the for range loop inside the goroutine to end when it completes a single iteration, the goroutine remains blocked at s.incoming <- evt.data within the loop. This prevents the goroutine from proceeding to the next iteration and exiting, thus causing the goroutine to remain in a blocked state, leading to a potential goroutine leak.

Proposed Solution:

To resolve this issue, I suggest implementing a solution where the child goroutine listens for both the done channel and the s.incoming channel, ensuring that the goroutine exits if either the done channel is closed or if there is a potential blockage on the s.incoming channel.

select {
case s.incoming <- evt.data:
case <-s.done:
    return
}

---

I am concerned that this could lead to unbounded goroutine accumulation, especially in high-load scenarios.

If I have misunderstood the behavior or missed any important context, I sincerely apologize for any confusion this might cause, and I would be grateful for further clarification. I’m happy to help address this issue if needed and contribute to improving the implementation.

Originally posted by @cryo-zd in #116

[jba]

@cryo-zd would you like to fix? If not, we would be happy to. Let us know.

[cryo-zd]

@cryo-zd would you like to fix? If not, we would be happy to. Let us know.

Thank you for the prompt and kind response!

I've implemented the suggested fix in PR #123 . I'd really appreciate your review and any feedback or corrections if I’ve missed something.

[samthanawalla]

It seems like #130 added a fix for this.

The new test is still very useful though!