Fix reliability findings

clients/go/mxgateway/session.go

@@ -418,7 +418,7 @@ func (s *Session) Events(ctx context.Context) (<-chan EventResult, error) {
 
 // EventsAfter streams ordered session events after the given worker sequence.
 func (s *Session) EventsAfter(ctx context.Context, afterWorkerSequence uint64) (<-chan EventResult, error) {
-	subscription, err := s.SubscribeEventsAfter(ctx, afterWorkerSequence)
+	subscription, err := s.subscribeEventsAfter(ctx, afterWorkerSequence, true)
 	if err != nil {
 		return nil, err
 	}
@@ -432,6 +432,10 @@ func (s *Session) SubscribeEvents(ctx context.Context) (*EventSubscription, erro
 
 // SubscribeEventsAfter starts an owned event subscription after the given worker sequence.
 func (s *Session) SubscribeEventsAfter(ctx context.Context, afterWorkerSequence uint64) (*EventSubscription, error) {
+	return s.subscribeEventsAfter(ctx, afterWorkerSequence, false)
+}
+
+func (s *Session) subscribeEventsAfter(ctx context.Context, afterWorkerSequence uint64, cancelWhenResultBufferFull bool) (*EventSubscription, error) {
 	streamCtx, cancel := context.WithCancel(ctx)
 	stream, err := s.client.StreamEventsRaw(streamCtx, &pb.StreamEventsRequest{
 		SessionId:           s.ID(),
@@ -450,7 +454,7 @@ func (s *Session) SubscribeEventsAfter(ctx context.Context, afterWorkerSequence
 		for {
 			event, err := stream.Recv()
 			if err == nil {
-				if !sendEventResult(streamCtx, results, EventResult{Event: event}) {
+				if !sendEventResult(streamCtx, results, EventResult{Event: event}, cancelWhenResultBufferFull, cancel) {
 					return
 				}
 				continue
@@ -458,7 +462,12 @@ func (s *Session) SubscribeEventsAfter(ctx context.Context, afterWorkerSequence
 			if err == io.EOF || status.Code(err) == codes.Canceled || streamCtx.Err() != nil {
 				return
 			}
-			sendEventResult(streamCtx, results, EventResult{Err: &GatewayError{Op: "stream events", Err: err}})
+			sendEventResult(
+				streamCtx,
+				results,
+				EventResult{Err: &GatewayError{Op: "stream events", Err: err}},
+				cancelWhenResultBufferFull,
+				cancel)
 			return
 		}
 	}()
@@ -477,7 +486,25 @@ func ensureBulkSize(name string, length int) error {
 	return nil
 }
 
-func sendEventResult(ctx context.Context, results chan<- EventResult, result EventResult) bool {
+func sendEventResult(
+	ctx context.Context,
+	results chan<- EventResult,
+	result EventResult,
+	cancelWhenBufferFull bool,
+	cancel context.CancelFunc,
+) bool {
+	if cancelWhenBufferFull {
+		select {
+		case results <- result:
+			return true
+		case <-ctx.Done():
+			return false
+		default:
+			cancel()
+			return false
+		}
+	}
+
 	select {
 	case results <- result:
 		return true