use Go 1.23 iterator

ecsta.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log/slog"
 	"os"
 	"strings"
 
@@ -80,12 +81,39 @@ func (app *Ecsta) describeTasks(ctx context.Context, opt *optionDescribeTasks) (
 	return out.Tasks, nil
 }
 
+func newAnyIterator[IN any, OUT any, OPT any](ctx context.Context, do func(context.Context, *IN, ...OPT) (*OUT, error), next func(*IN, *OUT) bool, in *IN, opt ...OPT) func(func(*OUT, error) bool) {
+	return func(yield func(*OUT, error) bool) {
+		for {
+			slog.Info("do", "in", in)
+			out, err := do(ctx, in, opt...)
+			if err != nil {
+				yield(nil, err)
+				return
+			}
+			if !yield(out, err) {
+				return
+			}
+			if !next(in, out) { // if next returns true, has a next page. continue
+				return
+			}
+		}
+	}
+}
+
+func ecsListTasksNext(in *ecs.ListTasksInput, out *ecs.ListTasksOutput) bool {
+	if out.NextToken == nil {
+		return false
+	}
+	in.NextToken = out.NextToken
+	return true
+}
+
 func (app *Ecsta) listTasks(ctx context.Context, opt *optionListTasks) ([]types.Task, error) {
 	tasks := []types.Task{}
 	var inputs []*ecs.ListTasksInput
 	if opt.family != nil && opt.service != nil {
 		// LisTasks does not support family and service at the same time
-		// spilit into two requests
+		// split into two requests
 		inputs = []*ecs.ListTasksInput{
 			{Cluster: &app.cluster, Family: opt.family},
 			{Cluster: &app.cluster, ServiceName: opt.service},
@@ -99,18 +127,16 @@ func (app *Ecsta) listTasks(ctx context.Context, opt *optionListTasks) ([]types.
 		for _, status := range []types.DesiredStatus{types.DesiredStatusRunning, types.DesiredStatusStopped} {
 			input := input
 			input.DesiredStatus = status
-			tp := ecs.NewListTasksPaginator(app.ecs, input)
-			for tp.HasMorePages() {
-				to, err := tp.NextPage(ctx)
+			for res, err := range newAnyIterator(ctx, app.ecs.ListTasks, ecsListTasksNext, input) {
 				if err != nil {
 					return nil, err
 				}
-				if len(to.TaskArns) == 0 {
+				if len(res.TaskArns) == 0 {
 					continue
 				}
 				out, err := app.ecs.DescribeTasks(ctx, &ecs.DescribeTasksInput{
 					Cluster: &app.cluster,
-					Tasks:   to.TaskArns,
+					Tasks:   res.TaskArns,
 					Include: []types.TaskField{"TAGS"},
 				})
 				if err != nil {

go.mod

@@ -1,6 +1,6 @@
 module github.com/fujiwara/ecsta
 
-go 1.22
+go 1.23
 
 toolchain go1.23.0
 

logs.go

@@ -124,8 +124,34 @@ func (o followOption) Follow() bool {
 	return o.follow || o.endTime.IsZero()
 }
 
+func (app *Ecsta) getLogEventsIterator(ctx context.Context, in *cloudwatchlogs.GetLogEventsInput) func(func(*cloudwatchlogs.GetLogEventsOutput, error) bool) {
+	return func(yield func(*cloudwatchlogs.GetLogEventsOutput, error) bool) {
+		for {
+			res, err := app.logs.GetLogEvents(ctx, in)
+			if err != nil {
+				var ne *logsTypes.ResourceNotFoundException
+				// log group or log stream not found, retry after 5 seconds
+				if errors.As(err, &ne) {
+					sleepWithContext(ctx, 5*time.Second)
+				} else {
+					yield(nil, err)
+				}
+				return
+			}
+			if !yield(res, nil) {
+				return
+			}
+			if res.NextForwardToken == nil {
+				return
+			}
+			in.NextToken = res.NextForwardToken
+			in.StartTime = nil
+			in.EndTime = nil
+		}
+	}
+}
+
 func (app *Ecsta) followLogs(ctx context.Context, opt *followOption) error {
-	var nextToken *string
 	in := &cloudwatchlogs.GetLogEventsInput{
 		LogGroupName:  &opt.logGroup,
 		LogStreamName: &opt.logStream,
@@ -135,52 +161,21 @@ func (app *Ecsta) followLogs(ctx context.Context, opt *followOption) error {
 	if !opt.Follow() {
 		in.EndTime = aws.Int64(timeToInt64msec(opt.endTime))
 	}
-
-FOLLOW:
-	for {
-		if err := sleepWithContext(ctx, time.Second); err != nil {
-			if errors.Is(err, context.Canceled) {
-				return nil
-			}
-			return err
-		}
-		if nextToken != nil {
-			in = &cloudwatchlogs.GetLogEventsInput{
-				LogGroupName:  &opt.logGroup,
-				LogStreamName: &opt.logStream,
-				Limit:         aws.Int32(1000),
-				NextToken:     nextToken,
-			}
-		}
-		res, err := app.logs.GetLogEvents(ctx, in)
+	for res, err := range app.getLogEventsIterator(ctx, in) {
 		if err != nil {
-			var ne *logsTypes.ResourceNotFoundException
-			// log group or log stream not found, retry after 5 seconds
-			if errors.As(err, &ne) {
-				sleepWithContext(ctx, 5*time.Second)
-			} else {
-				slog.Warn("failed to get log events", "error", err)
-			}
-			continue
+			return err
 		}
 		for _, e := range res.Events {
 			ts := msecToTime(aws.ToInt64(e.Timestamp))
 			if !opt.Follow() && ts.After(opt.endTime) {
-				break FOLLOW
+				break
 			}
 			fmt.Println(strings.Join([]string{
 				ts.Format(time.RFC3339Nano),
 				opt.containerName,
 				aws.ToString(e.Message),
 			}, "\t"))
 		}
-		if aws.ToString(nextToken) == aws.ToString(res.NextForwardToken) {
-			if !opt.Follow() {
-				break FOLLOW
-			}
-			continue
-		}
-		nextToken = res.NextForwardToken
 	}
 	return nil
 }