Improve gateway reliability and client e2e coverage

clients/go/mxgateway/session.go

@@ -22,6 +22,30 @@ type EventResult struct {
 	Err   error
 }
 
+// EventSubscription owns a running gateway event stream.
+type EventSubscription struct {
+	results <-chan EventResult
+	cancel  context.CancelFunc
+	done    <-chan struct{}
+	once    sync.Once
+}
+
+// Events returns the stream results channel.
+func (s *EventSubscription) Events() <-chan EventResult {
+	return s.results
+}
+
+// Close cancels the stream and waits for the receive goroutine to stop.
+func (s *EventSubscription) Close() {
+	if s == nil {
+		return
+	}
+	s.once.Do(func() {
+		s.cancel()
+		<-s.done
+	})
+}
+
 // Session represents one gateway-backed MXAccess session.
 type Session struct {
 	client     *Client
@@ -394,34 +418,56 @@ 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) {
-	stream, err := s.client.StreamEventsRaw(ctx, &pb.StreamEventsRequest{
+	subscription, err := s.SubscribeEventsAfter(ctx, afterWorkerSequence)
+	if err != nil {
+		return nil, err
+	}
+	return subscription.Events(), nil
+}
+
+// SubscribeEvents starts an owned event subscription.
+func (s *Session) SubscribeEvents(ctx context.Context) (*EventSubscription, error) {
+	return s.SubscribeEventsAfter(ctx, 0)
+}
+
+// SubscribeEventsAfter starts an owned event subscription after the given worker sequence.
+func (s *Session) SubscribeEventsAfter(ctx context.Context, afterWorkerSequence uint64) (*EventSubscription, error) {
+	streamCtx, cancel := context.WithCancel(ctx)
+	stream, err := s.client.StreamEventsRaw(streamCtx, &pb.StreamEventsRequest{
 		SessionId:           s.ID(),
 		AfterWorkerSequence: afterWorkerSequence,
 	})
 	if err != nil {
+		cancel()
 		return nil, err
 	}
 
 	results := make(chan EventResult, 16)
+	done := make(chan struct{})
 	go func() {
 		defer close(results)
+		defer close(done)
 		for {
 			event, err := stream.Recv()
 			if err == nil {
-				if !sendEventResult(ctx, results, EventResult{Event: event}) {
+				if !sendEventResult(streamCtx, results, EventResult{Event: event}) {
 					return
 				}
 				continue
 			}
-			if err == io.EOF || status.Code(err) == codes.Canceled || ctx.Err() != nil {
+			if err == io.EOF || status.Code(err) == codes.Canceled || streamCtx.Err() != nil {
 				return
 			}
-			sendEventResult(ctx, results, EventResult{Err: &GatewayError{Op: "stream events", Err: err}})
+			sendEventResult(streamCtx, results, EventResult{Err: &GatewayError{Op: "stream events", Err: err}})
 			return
 		}
 	}()
 
-	return results, nil
+	return &EventSubscription{
+		results: results,
+		cancel:  cancel,
+		done:    done,
+	}, nil
 }
 
 func ensureBulkSize(name string, length int) error {