Test AWS BYOC log stream handling

src/pkg/cli/client/byoc/aws/stream.go

@@ -17,6 +17,23 @@ import (
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
+type SSMode uint8
+
+const (
+	SSModeUnknown SSMode = iota
+	SSModeAwslogs
+	SSModeCd
+	SSModeFirelens
+	SSModeECS
+)
+
+type TailResponseHeader struct {
+	Service string
+	Host    string
+	Etag    string
+	Mode    SSMode
+}
+
 // byocServerStream is a wrapper around awsecs.EventStream that implements connect-like ServerStream
 type byocServerStream struct {
 	ctx      context.Context
@@ -58,15 +75,15 @@ func (bs *byocServerStream) Msg() *defangv1.TailResponse {
 }
 
 func (bs *byocServerStream) Receive() bool {
-	var evts []ecs.LogEvent
+	var events []ecs.LogEvent
 	select {
 	case e := <-bs.stream.Events(): // blocking
 		if bs.stream.Err() != nil {
 			bs.err = bs.stream.Err()
 			return false
 		}
 		var err error
-		evts, err = ecs.GetLogEvents(e)
+		events, err = ecs.GetLogEvents(e)
 		if err != nil {
 			bs.err = err
 			return false
@@ -76,106 +93,134 @@ func (bs *byocServerStream) Receive() bool {
 		return false
 	}
 
-	bs.response = bs.parseEvents(evts)
-
-	return true
-}
-
-func (bs *byocServerStream) parseEvents(events []ecs.LogEvent) *defangv1.TailResponse {
-	var response defangv1.TailResponse
 	if len(events) == 0 {
 		// The original gRPC/connect server stream would never send an empty response.
 		// We could loop around the select, but returning an empty response updates the spinner.
-		return nil
+		return true
 	}
-	parseFirelensRecords := false
-	parseECSEventRecords := false
+
 	// Get the Etag/Host/Service from the first event (should be the same for all events in this batch)
-	event := events[0]
-	if parts := strings.Split(*event.LogStreamName, "/"); len(parts) == 3 {
+	header := bs.makeResponseHeader(events[0])
+	// if unable to make a header, skip this stream
+	if header == nil {
+		return true
+	}
+
+	bs.response = bs.makeResponse(header, events)
+
+	return true
+}
+
+func (bs *byocServerStream) makeResponseHeader(event ecs.LogEvent) *TailResponseHeader {
+	var header TailResponseHeader
+
+	parts := strings.Split(*event.LogStreamName, "/")
+
+	if len(parts) == 3 {
 		if strings.Contains(*event.LogGroupIdentifier, ":"+byoc.CdTaskPrefix) {
 			// These events are from the CD task: "crun/main/taskID" stream; we should detect stdout/stderr
-			response.Etag = bs.etag // pass the etag filter below, but we already filtered the tail by taskID
-			response.Host = "pulumi"
-			response.Service = "cd"
-		} else {
-			// These events are from an awslogs service task: "tenant/service_etag/taskID" stream
-			response.Host = parts[2] // TODO: figure out actual hostname/IP
-			parts = strings.Split(parts[1], "_")
-			if len(parts) != 2 || !pkg.IsValidRandomID(parts[1]) {
-				// skip, ignore sidecar logs (like route53-sidecar or fluentbit)
-				return nil
-			}
-			service, etag := parts[0], parts[1]
-			response.Etag = etag
-			response.Service = service
+			header.Etag = bs.etag // pass the etag filter below, but we already filtered the tail by taskID
+			header.Host = "pulumi"
+			header.Service = "cd"
+			header.Mode = SSModeCd
+
+			return &header
+		}
+
+		// These events are from an awslogs service task with a stream name like "service/service_etag/taskID"
+		header.Host = parts[2] // TODO: figure out actual hostname/IP
+		seParts := strings.Split(parts[1], "_")
+		if len(seParts) != 2 || !pkg.IsValidRandomID(seParts[1]) {
+			// skip, ignore sidecar logs (like route53-sidecar or fluentbit)
+			return nil
 		}
-	} else if strings.Contains(*event.LogStreamName, "-firelens-") {
+
+		service, etag := seParts[0], seParts[1]
+		header.Etag = etag
+		header.Service = service
+		header.Mode = SSModeAwslogs
+
+		return &header
+	}
+
+	if strings.Contains(*event.LogStreamName, "-firelens-") {
 		// These events are from the Firelens sidecar; try to parse the JSON
 		var record logs.FirelensMessage
 		if err := json.Unmarshal([]byte(*event.Message), &record); err == nil {
-			response.Etag = record.Etag
-			response.Host = record.Host             // TODO: use "kaniko" for kaniko logs
-			response.Service = record.ContainerName // TODO: could be service_etag
-			parseFirelensRecords = true
+			header.Etag = record.Etag
+			header.Host = record.Host             // TODO: use "kaniko" for kaniko logs
+			header.Service = record.ContainerName // TODO: could be service_etag
+			header.Mode = SSModeFirelens
 		}
-	} else if strings.HasSuffix(*event.LogGroupIdentifier, "/ecs") || strings.HasSuffix(*event.LogGroupIdentifier, "/ecs:*") {
-		parseECSEventRecords = true
-		response.Etag = bs.etag
-		response.Service = "ecs"
-	}
 
-	// Client-side filtering
-	if bs.etag != "" && bs.etag != response.Etag {
-		return nil // TODO: filter these out using the AWS StartLiveTail API
+		return &header
 	}
 
-	if len(bs.services) > 0 && !pkg.Contains(bs.services, response.GetService()) {
-		return nil // TODO: filter these out using the AWS StartLiveTail API
+	if strings.HasSuffix(*event.LogGroupIdentifier, "/ecs") || strings.HasSuffix(*event.LogGroupIdentifier, "/ecs:*") {
+		header.Etag = bs.etag
+		header.Service = "ecs"
+		header.Mode = SSModeECS
+
+		return &header
 	}
 
+	return nil
+}
+
+func (bs *byocServerStream) makeResponse(header *TailResponseHeader, events []ecs.LogEvent) *defangv1.TailResponse {
 	entries := make([]*defangv1.LogEntry, 0, len(events))
+	response := &defangv1.TailResponse{
+		Service: header.Service,
+		Host:    header.Host,
+		Etag:    header.Etag,
+		Entries: entries,
+	}
+
 	for _, event := range events {
-		entry := &defangv1.LogEntry{
-			Message:   *event.Message,
-			Stderr:    false,
-			Timestamp: timestamppb.New(time.UnixMilli(*event.Timestamp)),
-		}
-		if parseFirelensRecords {
-			var record logs.FirelensMessage
-			if err := json.Unmarshal([]byte(entry.Message), &record); err == nil {
-				entry.Message = record.Log
-				if record.ContainerName == "kaniko" {
-					entry.Stderr = logs.IsLogrusError(entry.Message)
-				} else {
-					entry.Stderr = record.Source == logs.SourceStderr
-				}
-			}
-		} else if parseECSEventRecords {
-			evt, err := ecs.ParseECSEvent([]byte(*event.Message))
-			if err != nil {
-				term.Debugf("error parsing ECS event, output raw event log: %v", err)
+		entry := bs.makeEntry(header.Mode, event)
+		response.Entries = append(response.Entries, entry)
+	}
+
+	return response
+}
+
+func (bs *byocServerStream) makeEntry(mode SSMode, event ecs.LogEvent) *defangv1.LogEntry {
+	entry := defangv1.LogEntry{
+		Message:   *event.Message,
+		Stderr:    false,
+		Timestamp: timestamppb.New(time.UnixMilli(*event.Timestamp)),
+	}
+
+	switch mode {
+	case SSModeAwslogs:
+		var record logs.FirelensMessage
+		err := json.Unmarshal([]byte(entry.Message), &record)
+		if err != nil {
+			term.Debugf("error parsing Firelens message, output raw event log: %v", err)
+		} else {
+			entry.Message = record.Log
+			if record.ContainerName == "kaniko" {
+				entry.Stderr = logs.IsLogrusError(entry.Message)
 			} else {
-				bs.ecsEventsHandler.HandleECSEvent(evt)
-				entry.Service = evt.Service()
-				entry.Etag = evt.Etag()
-				entry.Host = evt.Host()
-				entry.Message = evt.Status()
+				entry.Stderr = record.Source == logs.SourceStderr
 			}
-		} else if response.Service == "cd" && strings.HasPrefix(entry.Message, logs.ErrorPrefix) {
-			entry.Stderr = true
 		}
-		if entry.Etag != "" && bs.etag != "" && entry.Etag != bs.etag {
-			continue
+	case SSModeECS:
+		evt, err := ecs.ParseECSEvent([]byte(*event.Message))
+		if err != nil {
+			term.Debugf("error parsing ECS event, output raw event log: %v", err)
+		} else {
+			bs.ecsEventsHandler.HandleECSEvent(evt)
+			entry.Service = evt.Service()
+			entry.Etag = evt.Etag()
+			entry.Host = evt.Host()
+			entry.Message = evt.Status()
 		}
-		if entry.Service != "" && len(bs.services) > 0 && !pkg.Contains(bs.services, entry.Service) {
-			continue
+	case SSModeCd:
+		if strings.HasPrefix(entry.Message, logs.ErrorPrefix) {
+			entry.Stderr = true
 		}
-		entries = append(entries, entry)
-	}
-	if len(entries) == 0 {
-		return nil
 	}
-	response.Entries = entries
-	return &response
+
+	return &entry
 }