backend: list message concurrent stream send fix

backend/pkg/api/connect/service/console/service.go

@@ -52,17 +52,6 @@ func NewService(logger *zap.Logger,
 
 // ListMessages consumes a Kafka topic and streams the Kafka records back.
 func (api *Service) ListMessages(ctx context.Context, req *connect.Request[v1alpha.ListMessagesRequest], stream *connect.ServerStream[v1alpha.ListMessagesResponse]) error {
-	timeout := 35 * time.Second
-	if req.Msg.GetFilterInterpreterCode() != "" || req.Msg.GetStartOffset() == console.StartOffsetNewest {
-		// Push-down filters and StartOffset = Newest may be long-running streams.
-		// There's already a client-side provided timeout which we usually trust.
-		// But additionally we want to ensure it never takes much longer than that.
-		timeout = 31 * time.Minute
-	}
-
-	ctx, cancel := context.WithTimeout(ctx, timeout)
-	defer cancel()
-
 	lmq := httptypes.ListMessagesRequest{
 		TopicName:             req.Msg.GetTopic(),
 		StartOffset:           req.Msg.GetStartOffset(),
@@ -122,6 +111,17 @@ func (api *Service) ListMessages(ctx context.Context, req *connect.Request[v1alp
 
 	api.authHooks.PrintListMessagesAuditLog(ctx, req, &listReq)
 
+	timeout := 35 * time.Second
+	if req.Msg.GetFilterInterpreterCode() != "" || req.Msg.GetStartOffset() == console.StartOffsetNewest {
+		// Push-down filters and StartOffset = Newest may be long-running streams.
+		// There's already a client-side provided timeout which we usually trust.
+		// But additionally we want to ensure it never takes much longer than that.
+		timeout = 31 * time.Minute
+	}
+
+	ctx, cancel := context.WithTimeoutCause(ctx, timeout, errors.New("list fetch timeout"))
+	defer cancel()
+
 	progress := &streamProgressReporter{
 		ctx:              ctx,
 		logger:           api.logger,

backend/pkg/api/connect/service/console/stream_progress_reporter.go

@@ -11,6 +11,7 @@ package console
 
 import (
 	"context"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -31,6 +32,8 @@ type streamProgressReporter struct {
 
 	messagesConsumed atomic.Int64
 	bytesConsumed    atomic.Int64
+
+	writeMutex sync.Mutex
 }
 
 func (p *streamProgressReporter) Start() {
@@ -44,41 +47,53 @@ func (p *streamProgressReporter) Start() {
 	// topic it may take some time until there are messages. This go routine is in charge of keeping the user up to
 	// date about the progress Kowl made streaming the topic
 	go func() {
+		ticker := time.NewTicker(1 * time.Second)
+
 		for {
 			select {
 			case <-p.ctx.Done():
+				ticker.Stop()
 				return
-			default:
+			case <-ticker.C:
 				p.reportProgress()
 			}
-			time.Sleep(1 * time.Second)
 		}
 	}()
 }
 
 func (p *streamProgressReporter) reportProgress() {
+	p.writeMutex.Lock()
+	defer p.writeMutex.Unlock()
+
 	msg := &v1alpha.ListMessagesResponse_ProgressMessage{
 		MessagesConsumed: p.messagesConsumed.Load(),
 		BytesConsumed:    p.bytesConsumed.Load(),
 	}
 
-	p.stream.Send(&v1alpha.ListMessagesResponse{
+	if err := p.stream.Send(&v1alpha.ListMessagesResponse{
 		ControlMessage: &v1alpha.ListMessagesResponse_Progress{
 			Progress: msg,
 		},
-	})
+	}); err != nil {
+		p.logger.Error("send error in stream reportProgress", zap.Error(err))
+	}
 }
 
 func (p *streamProgressReporter) OnPhase(name string) {
+	p.writeMutex.Lock()
+	defer p.writeMutex.Unlock()
+
 	msg := &v1alpha.ListMessagesResponse_PhaseMessage{
 		Phase: name,
 	}
 
-	p.stream.Send(&v1alpha.ListMessagesResponse{
+	if err := p.stream.Send(&v1alpha.ListMessagesResponse{
 		ControlMessage: &v1alpha.ListMessagesResponse_Phase{
 			Phase: msg,
 		},
-	})
+	}); err != nil {
+		p.logger.Error("send error in stream OnPhase", zap.Error(err))
+	}
 }
 
 func (p *streamProgressReporter) OnMessageConsumed(size int64) {
@@ -87,6 +102,13 @@ func (p *streamProgressReporter) OnMessageConsumed(size int64) {
 }
 
 func (p *streamProgressReporter) OnMessage(message *kafka.TopicMessage) {
+	if message == nil {
+		return
+	}
+
+	p.writeMutex.Lock()
+	defer p.writeMutex.Unlock()
+
 	headers := make([]*v1alpha.KafkaRecordHeader, 0, len(message.Headers))
 
 	for _, mh := range message.Headers {
@@ -164,36 +186,48 @@ func (p *streamProgressReporter) OnMessage(message *kafka.TopicMessage) {
 		})
 	}
 
-	p.stream.Send(&v1alpha.ListMessagesResponse{
+	if err := p.stream.Send(&v1alpha.ListMessagesResponse{
 		ControlMessage: &v1alpha.ListMessagesResponse_Data{
 			Data: data,
 		},
-	})
+	}); err != nil {
+		p.logger.Error("send error in stream OnMessage", zap.Error(err))
+	}
 }
 
 func (p *streamProgressReporter) OnComplete(elapsedMs int64, isCancelled bool) {
+	p.writeMutex.Lock()
+	defer p.writeMutex.Unlock()
+
 	msg := &v1alpha.ListMessagesResponse_StreamCompletedMessage{
 		ElapsedMs:        elapsedMs,
 		IsCancelled:      isCancelled,
 		MessagesConsumed: p.messagesConsumed.Load(),
 		BytesConsumed:    p.bytesConsumed.Load(),
 	}
 
-	p.stream.Send(&v1alpha.ListMessagesResponse{
+	if err := p.stream.Send(&v1alpha.ListMessagesResponse{
 		ControlMessage: &v1alpha.ListMessagesResponse_Done{
 			Done: msg,
 		},
-	})
+	}); err != nil {
+		p.logger.Error("send error in stream OnComplete", zap.Error(err))
+	}
 }
 
 func (p *streamProgressReporter) OnError(message string) {
+	p.writeMutex.Lock()
+	defer p.writeMutex.Unlock()
+
 	msg := &v1alpha.ListMessagesResponse_ErrorMessage{
 		Message: message,
 	}
 
-	p.stream.Send(&v1alpha.ListMessagesResponse{
+	if err := p.stream.Send(&v1alpha.ListMessagesResponse{
 		ControlMessage: &v1alpha.ListMessagesResponse_Error{
 			Error: msg,
 		},
-	})
+	}); err != nil {
+		p.logger.Error("send error in stream OnError", zap.Error(err))
+	}
 }

backend/pkg/kafka/consumer.go

@@ -115,8 +115,8 @@ func (s *Service) FetchMessages(ctx context.Context, progress IListMessagesProgr
 	// 2. Create consumer workers
 	jobs := make(chan *kgo.Record, 100)
 	resultsCh := make(chan *TopicMessage, 100)
-	workerCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
+	workerCtx, cancel := context.WithCancelCause(ctx)
+	defer cancel(errors.New("worker cancel"))
 
 	wg := sync.WaitGroup{}
 

frontend/src/state/backendApi.ts

@@ -245,6 +245,13 @@ const addBearerTokenInterceptor: ConnectRpcInterceptor = (next) => async (req: U
 };
 
 
+const transport = createConnectTransport({
+    baseUrl: appConfig.grpcBase,
+    interceptors: [addBearerTokenInterceptor]
+});
+
+const consoleClient = createPromiseClient(ConsoleService, transport);
+
 let messageSearchAbortController: AbortController | null = null;
 
 //
@@ -372,14 +379,9 @@ const apiStore = {
 
         // do it
         const abortController = messageSearchAbortController = new AbortController();
-        const transport = createConnectTransport({
-            baseUrl: appConfig.grpcBase,
-            interceptors: [addBearerTokenInterceptor]
-        });
-
-        const client = createPromiseClient(ConsoleService, transport);
 
         const req = new ListMessagesRequest();
+
         req.topic = searchRequest.topicName;
         req.startOffset = BigInt(searchRequest.startOffset);
         req.startTimestamp = BigInt(searchRequest.startTimestamp);
@@ -398,8 +400,12 @@ const apiStore = {
             timeoutMs = 30 * minuteMs;
         }
 
+        abortController.signal.addEventListener('abort', () => {
+            console.log('ABORTED');
+        })
+
         try {
-            for await (const res of await client.listMessages(req, { signal: abortController.signal, timeoutMs })) {
+            for await (const res of await consoleClient.listMessages(req, { signal: abortController.signal, timeoutMs })) {
                 if (abortController.signal.aborted)
                     break;
 
@@ -415,6 +421,7 @@ const apiStore = {
                             this.messagesTotalConsumed = Number(res.controlMessage.value.messagesConsumed);
                             break;
                         case 'done':
+                            console.log('done phase')
                             this.messagesElapsedMs = Number(res.controlMessage.value.elapsedMs);
                             this.messagesBytesConsumed = Number(res.controlMessage.value.bytesConsumed);
                             // this.MessageSearchCancelled = msg.isCancelled;