Implement Subscribe API call for BYOC (#613)

* Implement Subscribe API call for BYOC

also fix collection stream error handling

* Add subscribe test and output whole message in one call in tail

* Remove Errs() method from collection stream

Keep all channel logic inside collection stream code

* Abstract ECS event parsing

* Fix event filtering, use defer to make sure term clean up

* Restructure tail log message construction

* Simplify append

Co-authored-by: Lio李歐 <[email protected]>

* Revert "Simplify append"

This reverts commit ba38ee9.

---------

Co-authored-by: Edward J <[email protected]>
Co-authored-by: Lio李歐 <[email protected]>

Author: 3 people

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

@@ -12,7 +12,7 @@ import (
 	"os"
 	"sort"
 	"strings"
-	"time"
+	"sync"
 
 	"github.com/DefangLabs/defang/src/pkg"
 	"github.com/DefangLabs/defang/src/pkg/cli/client"
@@ -41,6 +41,9 @@ type ByocAws struct {
 	cdTasks      map[string]ecs.TaskArn
 	driver       *cfn.AwsEcs
 	publicNatIps []string
+
+	ecsEventHandlers []ECSEventHandler
+	handlersLock     sync.RWMutex
 }
 
 var _ client.Client = (*ByocAws)(nil)
@@ -452,8 +455,6 @@ func (b *ByocAws) Follow(ctx context.Context, req *defangv1.TailRequest) (client
 		return nil, err
 	}
 
-	ctx, cancel := context.WithCancelCause(ctx)
-
 	etag := req.Etag
 	// if etag == "" && req.Service == "cd" {
 	// 	etag = awsecs.GetTaskID(b.cdTaskArn); TODO: find the last CD task
@@ -493,16 +494,7 @@ func (b *ByocAws) Follow(ctx context.Context, req *defangv1.TailRequest) (client
 		return nil, annotateAwsError(err)
 	}
 
-	if taskArn != nil {
-		go func() {
-			if err := ecs.WaitForTask(ctx, taskArn, 3*time.Second); err != nil {
-				time.Sleep(time.Second) // make sure we got all the logs from the task before cancelling
-				cancel(err)
-			}
-		}()
-	}
-
-	return newByocServerStream(ctx, eventStream, etag, req.GetServices()), nil
+	return newByocServerStream(ctx, eventStream, etag, req.GetServices(), b), nil
 }
 
 // This function was copied from Fabric controller and slightly modified to work with BYOC
@@ -781,6 +773,32 @@ func ensure(cond bool, msg string) {
 	}
 }
 
-func (b *ByocAws) Subscribe(context.Context, *defangv1.SubscribeRequest) (client.ServerStream[defangv1.SubscribeResponse], error) {
-	return nil, client.ErrNotImplemented("not yet implemented for BYOC; please use the AWS ECS dashboard") // FIXME: implement this for BYOC
+type ECSEventHandler interface {
+	HandleECSEvent(evt ecs.Event)
+}
+
+func (b *ByocAws) Subscribe(ctx context.Context, req *defangv1.SubscribeRequest) (client.ServerStream[defangv1.SubscribeResponse], error) {
+	s := &byocSubscribeServerStream{
+		services: req.Services,
+		etag:     req.Etag,
+		ctx:      ctx,
+
+		ch: make(chan *defangv1.SubscribeResponse),
+	}
+	b.AddEcsEventHandler(s)
+	return s, nil
+}
+
+func (b *ByocAws) HandleECSEvent(evt ecs.Event) {
+	b.handlersLock.RLock()
+	defer b.handlersLock.RUnlock()
+	for _, handler := range b.ecsEventHandlers {
+		handler.HandleECSEvent(evt)
+	}
+}
+
+func (b *ByocAws) AddEcsEventHandler(handler ECSEventHandler) {
+	b.handlersLock.Lock()
+	defer b.handlersLock.Unlock()
+	b.ecsEventHandlers = append(b.ecsEventHandlers, handler)
 }

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

@@ -1,10 +1,19 @@
 package aws
 
 import (
+	"bufio"
+	"bytes"
 	"context"
+	"embed"
+	"encoding/json"
+	"io"
+	"path"
+	"strings"
+	"sync"
 	"testing"
 
 	"github.com/DefangLabs/defang/src/pkg/cli/client"
+	"github.com/DefangLabs/defang/src/pkg/clouds/aws/ecs"
 	"github.com/DefangLabs/defang/src/pkg/types"
 	defangv1 "github.com/DefangLabs/defang/src/protos/io/defang/v1"
 	compose "github.com/compose-spec/compose-go/v2/types"
@@ -74,3 +83,83 @@ func (f FakeLoader) LoadProject(ctx context.Context) (*compose.Project, error) {
 func (f FakeLoader) LoadProjectName(ctx context.Context) (string, error) {
 	return f.ProjectName, nil
 }
+
+//go:embed test_ecs_events/*.json
+var testDir embed.FS
+
+//go:embed test_ecs_events/*.events
+var expectedDir embed.FS
+
+func TestSubscribe(t *testing.T) {
+	t.Skip("Pending test")
+	tests, err := testDir.ReadDir("test_ecs_events")
+	if err != nil {
+		t.Fatalf("failed to load ecs events test files: %v", err)
+	}
+	for _, tt := range tests {
+		t.Run(tt.Name(), func(t *testing.T) {
+			start := strings.LastIndex(tt.Name(), "-")
+			end := strings.LastIndex(tt.Name(), ".")
+			if start == -1 || end == -1 {
+				t.Fatalf("cannot find etag from invalid test file name: %s", tt.Name())
+			}
+			name := tt.Name()[:start]
+			etag := tt.Name()[start+1 : end]
+
+			byoc := &ByocAws{}
+
+			resp, err := byoc.Subscribe(context.Background(), &defangv1.SubscribeRequest{
+				Etag:     etag,
+				Services: []string{"api", "web"},
+			})
+			if err != nil {
+				t.Fatalf("Subscribe() failed: %v", err)
+			}
+
+			var wg sync.WaitGroup
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+
+				filename := path.Join("test_ecs_events", name+".events")
+				ef, _ := expectedDir.ReadFile(filename)
+				dec := json.NewDecoder(bytes.NewReader(ef))
+
+				for {
+					if !resp.Receive() {
+						if resp.Err() != nil {
+							t.Errorf("Receive() failed: %v", resp.Err())
+						}
+						break
+					}
+					msg := resp.Msg()
+					var expected defangv1.SubscribeResponse
+					if err := dec.Decode(&expected); err == io.EOF {
+						t.Errorf("unexpected message: %v", msg)
+					} else if err != nil {
+						t.Errorf("error unmarshaling expected ECS event: %v", err)
+					} else if msg.Name != expected.Name || msg.Status != expected.Status || msg.State != expected.State {
+						t.Errorf("expected message-, got+\n-%v\n+%v", &expected, msg)
+					}
+				}
+			}()
+
+			data, err := testDir.ReadFile(path.Join("test_ecs_events", tt.Name()))
+			if err != nil {
+				t.Fatalf("failed to read test file: %v", err)
+			}
+			lines := bufio.NewScanner(bytes.NewReader(data))
+			for lines.Scan() {
+				ecsEvt, err := ecs.ParseECSEvent([]byte(lines.Text()))
+				if err != nil {
+					t.Fatalf("error parsing ECS event: %v", err)
+				}
+
+				byoc.HandleECSEvent(ecsEvt)
+			}
+			resp.Close()
+
+			wg.Wait()
+		})
+	}
+}

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

@@ -3,9 +3,7 @@ package aws
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 	"io"
-	"path"
 	"strings"
 	"time"
 
@@ -23,25 +21,22 @@ import (
 type byocServerStream struct {
 	ctx      context.Context
 	err      error
-	errCh    <-chan error
 	etag     string
 	response *defangv1.TailResponse
 	services []string
 	stream   ecs.EventStream
-}
 
-func newByocServerStream(ctx context.Context, stream ecs.EventStream, etag string, services []string) *byocServerStream {
-	var errCh <-chan error
-	if errch, ok := stream.(hasErrCh); ok {
-		errCh = errch.Errs()
-	}
+	ecsEventsHandler ECSEventHandler
+}
 
+func newByocServerStream(ctx context.Context, stream ecs.EventStream, etag string, services []string, ecsEventHandler ECSEventHandler) *byocServerStream {
 	return &byocServerStream{
 		ctx:      ctx,
-		errCh:    errCh,
 		etag:     etag,
 		stream:   stream,
 		services: services,
+
+		ecsEventsHandler: ecsEventHandler,
 	}
 }
 
@@ -62,24 +57,20 @@ func (bs *byocServerStream) Msg() *defangv1.TailResponse {
 	return bs.response
 }
 
-type hasErrCh interface {
-	Errs() <-chan error
-}
-
 func (bs *byocServerStream) Receive() bool {
 	var evts []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)
 		if err != nil {
 			bs.err = err
 			return false
 		}
-	case err := <-bs.errCh: // blocking (if not nil)
-		bs.err = err
-		return false // abort on first error?
-
 	case <-bs.ctx.Done(): // blocking (if not nil)
 		bs.err = context.Cause(bs.ctx)
 		return false
@@ -132,19 +123,21 @@ func (bs *byocServerStream) parseEvents(events []ecs.LogEvent) (*defangv1.TailRe
 		}
 	} 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, nil // TODO: filter these out using the AWS StartLiveTail API
 	}
 
-	if len(bs.services) > 0 && !pkg.Contains(bs.services, bs.response.GetService()) {
+	if len(bs.services) > 0 && !pkg.Contains(bs.services, response.GetService()) {
 		return nil, nil // TODO: filter these out using the AWS StartLiveTail API
 	}
 
-	entries := make([]*defangv1.LogEntry, len(events))
-	for i, event := range events {
+	entries := make([]*defangv1.LogEntry, 0, len(events))
+	for _, event := range events {
 		entry := &defangv1.LogEntry{
 			Message:   *event.Message,
 			Stderr:    false,
@@ -162,81 +155,36 @@ func (bs *byocServerStream) parseEvents(events []ecs.LogEvent) (*defangv1.TailRe
 			}
 		} else if parseECSEventRecords {
 			var err error
-			if err = parseECSEventRecord(event, entry); err != nil {
+			if err = bs.parseECSEventRecord(event, entry); err != nil {
 				term.Debugf("error parsing ECS event, output raw event log: %v", err)
 			}
 		} else if response.Service == "cd" && strings.HasPrefix(entry.Message, " ** ") {
 			entry.Stderr = true
 		}
-		entries[i] = entry
+		if entry.Etag != "" && bs.etag != "" && entry.Etag != bs.etag {
+			continue
+		}
+		if entry.Service != "" && len(bs.services) > 0 && !pkg.Contains(bs.services, entry.Service) {
+			continue
+		}
+		entries = append(entries, entry)
+	}
+	if len(entries) == 0 {
+		return nil, nil
 	}
 	response.Entries = entries
 	return &response, nil
 }
 
-func parseECSEventRecord(event ecs.LogEvent, entry *defangv1.LogEntry) error {
-	var ecsEvt ecs.Event
-	if err := json.Unmarshal([]byte(*event.Message), &ecsEvt); err != nil {
-		return fmt.Errorf("error unmarshaling ECS event: %w", err)
-	}
-
-	var buf strings.Builder
-	fmt.Fprintf(&buf, "%s ", ecsEvt.DetailType)
-	if len(ecsEvt.Resources) > 0 {
-		fmt.Fprintf(&buf, "%s ", path.Base(ecsEvt.Resources[0]))
-	}
-	switch ecsEvt.DetailType {
-	case "ECS Task State Change":
-		var detail ecs.ECSTaskStateChange
-		if err := json.Unmarshal(ecsEvt.Detail, &detail); err != nil {
-			return fmt.Errorf("error unmarshaling ECS task state change: %w", err)
-		}
-
-		// Container name is in the format of "service_etag"
-		if len(detail.Containers) < 1 {
-			return fmt.Errorf("error parsing ECS task state change: missing containers section")
-		}
-		i := strings.LastIndex(detail.Containers[0].Name, "_")
-		if i < 0 {
-			return fmt.Errorf("error parsing ECS task state change: invalid container name %q", detail.Containers[0].Name)
-		}
-		entry.Service = detail.Containers[0].Name[:i]
-		entry.Etag = detail.Containers[0].Name[i+1:]
-		entry.Host = path.Base(ecsEvt.Resources[0])
-		fmt.Fprintf(&buf, "%s %s", path.Base(detail.ClusterArn), detail.LastStatus)
-		if detail.StoppedReason != "" {
-			fmt.Fprintf(&buf, " : %s", detail.StoppedReason)
-		}
-	case "ECS Service Action", "ECS Deployment State Change": // pretty much the same JSON structure for both
-		var detail ecs.ECSDeploymentStateChange
-		if err := json.Unmarshal(ecsEvt.Detail, &detail); err != nil {
-			return fmt.Errorf("error unmarshaling ECS service/deployment event: %v", err)
-		}
-		ecsSvcName := path.Base(ecsEvt.Resources[0])
-		// TODO: etag is not available at service and deployment level, find a possible correlation, possibly task definition revision using the deploymentId
-		snStart := strings.LastIndex(ecsSvcName, "_") // ecsSvcName is in the format "project_service-random", our validation does not allow '_' in service names
-		snEnd := strings.LastIndex(ecsSvcName, "-")
-		if snStart < 0 || snEnd < 0 || snStart >= snEnd {
-			return fmt.Errorf("error parsing ECS service action: invalid service name %q", ecsEvt.Resources[0])
-		}
-		entry.Service = ecsSvcName[snStart+1 : snEnd]
-		entry.Host = detail.DeploymentId
-		fmt.Fprintf(&buf, "%s", detail.EventName)
-		if detail.Reason != "" {
-			fmt.Fprintf(&buf, " : %s", detail.Reason)
-		}
-	default:
-		entry.Service = "ecs"
-		if len(ecsEvt.Resources) > 0 {
-			entry.Host = path.Base(ecsEvt.Resources[0])
-		}
-		// Print the unrecogonalized ECS event detail in prettry JSON format if possible
-		raw, err := json.MarshalIndent(ecsEvt.Detail, "", "  ")
-		if err != nil {
-			raw = []byte(ecsEvt.Detail)
-		}
-		fmt.Fprintf(&buf, "\n%s", raw)
+func (bs *byocServerStream) parseECSEventRecord(event ecs.LogEvent, entry *defangv1.LogEntry) error {
+	evt, err := ecs.ParseECSEvent([]byte(*event.Message))
+	if err != nil {
+		return err
 	}
-	entry.Message = buf.String()
+	bs.ecsEventsHandler.HandleECSEvent(evt)
+	entry.Service = evt.Service()
+	entry.Etag = evt.Etag()
+	entry.Host = evt.Host()
+	entry.Message = evt.Status()
 	return nil
 }