Please add context.Context support for RecvMsg()

[GiedriusS]

Use case(s) - what problem will this feature solve?

We want to write a high-performance loop that receives messages from a stream as fast as possible whilst at the same time protecting ourselves from one operation taking too long i.e. we want to return a partial response to our users if one stream is taking too long.

Proposed Solution

RecvMsg() needs to accept a context.Context.

Alternatives Considered

The proposed alternative to spawn a goroutine for each RecvMsg() seems horrible for performance. Adding context.Context support seems straightforward and it isn't such a huge breaking change for users to simply pass context.Background().

Additional Context

#445 it's really not a solution to spawn new goroutines for each RecvMsg() if you are receiving hundreds of thousands of messages potentially.

[dfawley]

This is no different from a typical io.Reader, which doesn't accept a context, either. You shouldn't need more than one goroutine for all of your reads in this case. Consider:

func readFromStream(stream pb.Service_MyHandlerServer, ch chan pb.MyMessageType) {
	defer close(ch)
	for {
		res, err := stream.RecvMsg()
		if err != nil {
			// error handling
			return
		}
		select {
		case ch<-res:
		case <-stream.Context().Done():
			// In case the receiver from the channel stops and cancels the RPC.
			// Don't do anything here and let the stream return an error.
		}
	}
}

func myClient() {
	// ...
	ch := make(chan pb.MyMessageType)
	go readFromStream(stream, ch)
	for {
		ctx, cancel := context.WithTimeout(ctx, time.Second)
		// ensure cancel() is called after this loop iteration
		select {
		case msg:=<-ch:
			// Process msg
		case <-ctx.Done():
			// Process timeout
			streamCtxCancel() // unblocks readFromStream by canceling the RPC
		}
	}
}

[github-actions]

This issue is labeled as requiring an update from the reporter, and no update has been received after 6 days. If no update is provided in the next 7 days, this issue will be automatically closed.

[GiedriusS]

Unfortunately but using a separate goroutine isn't a good solution performance-wise. There is a huge reduction in performance when you try to pass tens of thousands of received messages via the channel. In our case doing:

for {
  res, err := stream.RecvMsg()
  ...
}

vs. doing that via a channel reduces the query duration by a few seconds (thanos-io/thanos#5296). That's why I have created this issue (: perhaps there is some other solution?

[dfawley]

Other ideas which avoid the channel:

Using a per-operation timer:

ctx, cancel := context.WithCancel(ctx)
defer cancel()
stream, err := client.StreamingRPC(ctx)

for {
	t := time.AfterFunc(2*time.Second, cancel)
	res, err := stream.RecvMsg()
	t.Stop()
	// Handle res/err inline as before
}

Using a monitoring goroutine to cancel the stream context when progress stops (may not work if stream response processing is slow/unpredictable):

ctx, cancel := context.WithCancel(ctx)
defer cancel()
stream, err := client.StreamingRPC(ctx)

// Monitor messages received and cancel the stream if progress
// isn't made every two seconds.
msgs := int32(0)
go func() {
	lastMsgs := int32(0)
	for ctx.Err() == nil {
		time.Sleep(2*time.Seconds)
		t := atomic.LoadInt32(msgs)
		if lastMsgs == t {
			cancel()
			return
		}
		lastMsgs = t
	}
}()

for {
	res, err := stream.RecvMsg()
	atomic.AddInt32(msgs,1)
	// Handle res/err inline as before
}