Make limit int32 (#1609)

Author: lionello

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

@@ -686,7 +686,7 @@ func (b *ByocAws) QueryLogs(ctx context.Context, req *defangv1.TailRequest) (cli
 				etag = "" // no need to filter by etag
 			}
 		} else {
-			tailStream, err = b.driver.QueryTaskID(ctx, etag, time.Time{}, time.Now(), int(req.Limit))
+			tailStream, err = b.driver.QueryTaskID(ctx, etag, time.Time{}, time.Now(), req.Limit)
 			if err == nil {
 				b.cdTaskArn, err = b.driver.GetTaskArn(etag)
 				etag = "" // no need to filter by etag
@@ -713,12 +713,11 @@ func (b *ByocAws) QueryLogs(ctx context.Context, req *defangv1.TailRequest) (cli
 				lgis...,
 			)
 		} else {
-			limit := int(req.Limit)
 			evtsChan, errsChan := ecs.QueryLogGroups(
 				ctx,
 				start,
 				end,
-				limit,
+				req.Limit,
 				lgis...,
 			)
 			if evtsChan == nil {

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

@@ -561,7 +561,7 @@ func (b *ByocGcp) QueryLogs(ctx context.Context, req *defangv1.TailRequest) (cli
 		if req.Follow {
 			subscribeStream.StartFollow(since)
 		} else {
-			subscribeStream.Start(int(req.Limit))
+			subscribeStream.Start(req.Limit)
 		}
 
 		var cancel context.CancelCauseFunc
@@ -624,7 +624,7 @@ func (b *ByocGcp) QueryLogs(ctx context.Context, req *defangv1.TailRequest) (cli
 	if req.Follow {
 		logStream.StartFollow(startTime)
 	} else {
-		logStream.Start(int(req.Limit))
+		logStream.Start(req.Limit)
 	}
 	return logStream, nil
 }

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

@@ -125,7 +125,7 @@ func (s *ServerStream[T]) StartFollow(start time.Time) {
 	}()
 }
 
-func (s *ServerStream[T]) Start(limit int) {
+func (s *ServerStream[T]) Start(limit int32) {
 	query := s.query.GetQuery()
 	term.Debugf("Query logs with query: \n%v", query)
 	go func() {
@@ -146,7 +146,7 @@ func (s *ServerStream[T]) queryHead(query string) {
 	}
 }
 
-func (s *ServerStream[T]) queryTail(query string, limit int) {
+func (s *ServerStream[T]) queryTail(query string, limit int32) {
 	lister, err := s.gcp.ListLogEntries(s.ctx, query, gcp.OrderDescending)
 	if err != nil {
 		s.errCh <- err
@@ -171,7 +171,7 @@ func (s *ServerStream[T]) queryTail(query string, limit int) {
 	}
 }
 
-func (s *ServerStream[T]) listToBuffer(lister *gcp.Lister, limit int) ([]*T, error) {
+func (s *ServerStream[T]) listToBuffer(lister *gcp.Lister, limit int32) ([]*T, error) {
 	received := 0
 	buffer := make([]*T, 0, limit)
 	for range limit {
@@ -191,7 +191,7 @@ func (s *ServerStream[T]) listToBuffer(lister *gcp.Lister, limit int) ([]*T, err
 	return buffer, nil
 }
 
-func (s *ServerStream[T]) listToChannel(lister *gcp.Lister, limit int) error {
+func (s *ServerStream[T]) listToChannel(lister *gcp.Lister, limit int32) error {
 	received := 0
 	for {
 		entry, err := lister.Next()
@@ -206,7 +206,7 @@ func (s *ServerStream[T]) listToChannel(lister *gcp.Lister, limit int) error {
 			s.respCh <- resp
 		}
 		received += len(resps)
-		if limit > 0 && received >= limit {
+		if limit > 0 && received >= int(limit) {
 			return io.EOF
 		}
 	}

src/pkg/clouds/aws/ecs/logs.go

@@ -124,16 +124,15 @@ func TailLogGroup(ctx context.Context, input LogGroupInput) (LiveTailStream, err
 	return slto.GetStream(), nil
 }
 
-func QueryLogGroups(ctx context.Context, start, end time.Time, limit int, logGroups ...LogGroupInput) (<-chan LogEvent, <-chan error) {
-	// Gather logs from the CD task, kaniko, ECS events, and all services
+func QueryLogGroups(ctx context.Context, start, end time.Time, limit int32, logGroups ...LogGroupInput) (<-chan LogEvent, <-chan error) {
 	var evtsChan chan LogEvent
 	var errChan chan error
 	for _, lgi := range logGroups {
 		lgEvtChan := make(chan LogEvent)
 		// Start a go routine for each log group
 		go func(lgi LogGroupInput) {
 			defer close(lgEvtChan)
-			if err := QueryLogGroup(ctx, lgi, start, end, func(logEvents []LogEvent) error {
+			if err := QueryLogGroup(ctx, lgi, start, end, limit, func(logEvents []LogEvent) error {
 				for _, event := range logEvents {
 					lgEvtChan <- event
 				}
@@ -145,22 +144,22 @@ func QueryLogGroups(ctx context.Context, start, end time.Time, limit int, logGro
 		evtsChan = mergeLogEventChan(evtsChan, lgEvtChan) // Merge sort the log events based on timestamp
 		// take the last n events only
 		if limit > 0 {
-			evtsChan = takeLastN(evtsChan, limit)
+			evtsChan = takeLastN(evtsChan, int(limit))
 		}
 	}
 	return evtsChan, errChan
 }
 
-func QueryLogGroup(ctx context.Context, input LogGroupInput, start, end time.Time, cb func([]LogEvent) error) error {
+func QueryLogGroup(ctx context.Context, input LogGroupInput, start, end time.Time, limit int32, cb func([]LogEvent) error) error {
 	region := region.FromArn(input.LogGroupARN)
 	cw, err := newCloudWatchLogsClient(ctx, region)
 	if err != nil {
 		return err
 	}
-	return filterLogEvents(ctx, cw, input, start, end, cb)
+	return filterLogEvents(ctx, cw, input, start, end, limit, cb)
 }
 
-func filterLogEvents(ctx context.Context, cw *cloudwatchlogs.Client, lgi LogGroupInput, start, end time.Time, cb func([]LogEvent) error) error {
+func filterLogEvents(ctx context.Context, cw *cloudwatchlogs.Client, lgi LogGroupInput, start, end time.Time, limit int32, cb func([]LogEvent) error) error {
 	var pattern *string
 	if lgi.LogEventFilterPattern != "" {
 		pattern = &lgi.LogEventFilterPattern
@@ -189,6 +188,11 @@ func filterLogEvents(ctx context.Context, cw *cloudwatchlogs.Client, lgi LogGrou
 		params.LogStreamNamePrefix = &lgi.LogStreamNamePrefix
 	}
 	for {
+		if limit > 0 {
+			// Specifying the limit parameter only guarantees that a single page doesn't return more log events than the
+			// specified limit, but it might return fewer events than the limit. This is the expected API behavior.
+			params.Limit = ptr.Int32(limit)
+		}
 		fleo, err := cw.FilterLogEvents(ctx, params)
 		if err != nil {
 			return err
@@ -209,6 +213,10 @@ func filterLogEvents(ctx context.Context, cw *cloudwatchlogs.Client, lgi LogGrou
 		if fleo.NextToken == nil {
 			return nil
 		}
+		if limit > 0 && len(events) >= int(limit) {
+			return nil
+		}
+		limit -= int32(len(events)) // #nosec G115 - always safe because len(events) <= limit
 		params.NextToken = fleo.NextToken
 	}
 }

src/pkg/clouds/aws/ecs/stream.go

@@ -51,7 +51,7 @@ func QueryAndTailLogGroup(ctx context.Context, lgi LogGroupInput, start, end tim
 				end = time.Now()
 			}
 			// Query the logs between the start time and now; TODO: could use a single CloudWatch client for all queries in same region
-			if err := QueryLogGroup(ctx, lgi, start, end, func(events []LogEvent) error {
+			if err := QueryLogGroup(ctx, lgi, start, end, 0, func(events []LogEvent) error {
 				es.ch <- &types.StartLiveTailResponseStreamMemberSessionUpdate{
 					Value: types.LiveTailSessionUpdate{SessionResults: events},
 				}

src/pkg/clouds/aws/ecs/tail.go

@@ -57,7 +57,7 @@ func (a *AwsEcs) GetTaskArn(taskID string) (TaskArn, error) {
 	return &taskArn, nil
 }
 
-func (a *AwsEcs) QueryTaskID(ctx context.Context, taskID string, start, end time.Time, limit int) (EventStream[types.StartLiveTailResponseStream], error) {
+func (a *AwsEcs) QueryTaskID(ctx context.Context, taskID string, start, end time.Time, limit int32) (EventStream[types.StartLiveTailResponseStream], error) {
 	if taskID == "" {
 		return nil, errors.New("taskID is empty")
 	}
@@ -68,7 +68,8 @@ func (a *AwsEcs) QueryTaskID(ctx context.Context, taskID string, start, end time
 	}
 
 	lgi := LogGroupInput{LogGroupARN: a.LogGroupARN, LogStreamNames: []string{GetCDLogStreamForTaskID(taskID)}}
-	if err := QueryLogGroup(ctx, lgi, start, end, func(events []LogEvent) error {
+	// Note: this function only returns once the query is complete, so returning an event stream is somewhat misleading
+	if err := QueryLogGroup(ctx, lgi, start, end, limit, func(events []LogEvent) error {
 		es.ch <- &types.StartLiveTailResponseStreamMemberSessionUpdate{
 			Value: types.LiveTailSessionUpdate{SessionResults: events},
 		}
@@ -77,10 +78,6 @@ func (a *AwsEcs) QueryTaskID(ctx context.Context, taskID string, start, end time
 		es.err = err
 	}
 
-	if limit > 0 {
-		es.ch = takeLastN(es.ch, limit)
-	}
-
 	return es, nil
 }