Make batch run --follow actually stream, and terminate

api/client.go

@@ -585,35 +585,27 @@ func (client *Client) CancelBatchExecution(project, service, execution string) e
 	return client.request(http.MethodPost, path, &resp, nil)
 }
 
-// StreamBatchExecutionLogs writes one execution's container output to w.
+// ReadBatchExecutionLogs writes an execution's output from byteOffset onward and
+// returns how many bytes it wrote.
 //
 // Batch logs are addressed per execution rather than per service: a batch
 // service has no deployment for the service-wide endpoint to derive a label
 // from.
 //
-// A followed stream reconnects and resumes at an exact byte offset, the same
-// way service logs do. It always starts from the beginning of the log, since an
-// execution's output is finite and usually short -- there is no reason to
-// anchor to the end and risk missing the first lines.
-func (client *Client) StreamBatchExecutionLogs(project, service, execution string, w io.Writer, follow bool) error {
-	base := batchPath(project, service) + "/" + execution + "/logs"
-	logPath := func(params url.Values) string {
-		if follow {
-			params.Set("follow", "true")
-			params.Set("keepalive", keepaliveInterval.String())
-		}
-		return base + "?" + params.Encode()
-	}
-	start := url.Values{"seek": {"start"}}
-	if !follow {
-		return client.streamOnce(logPath(start), w)
-	}
-	return client.follow(followOptions{path: func(resumeAfter int) string {
-		if resumeAfter == 0 {
-			return logPath(start)
-		}
-		return logPath(url.Values{"seek": {"start"}, "offset": {strconv.Itoa(resumeAfter)}})
-	}}, w)
+// This is a plain read rather than a server-side follow. A followed stream is
+// never closed by the server when an execution's log completes, so a follower
+// has no way to learn it is done and hangs. Reading from an offset lets the
+// caller decide when to stop -- which for a batch execution is when the
+// execution itself reaches a terminal state.
+func (client *Client) ReadBatchExecutionLogs(project, service, execution string, byteOffset int, w io.Writer) (int, error) {
+	params := url.Values{"seek": {"start"}}
+	if byteOffset > 0 {
+		params.Set("offset", strconv.Itoa(byteOffset))
+	}
+	path := batchPath(project, service) + "/" + execution + "/logs?" + params.Encode()
+	counter := &countingWriter{w: w}
+	err := client.streamOnce(path, counter)
+	return counter.bytes, err
 }
 
 // GetBatchConfig retrieves a batch service's configuration.

cmd/batch.go

@@ -57,20 +57,15 @@ latest successful build is used.`,
 			if !batchRunFollow && !batchRunWait {
 				return nil
 			}
-			final, err := awaitBatchCompletion(client, cfg, exec.ID)
-			if err != nil {
-				return err
-			}
 			if batchRunFollow {
-				// Deliberately fetched after completion rather than streamed
-				// live. Batch containers routinely finish in under a second, so
-				// a tail started on "running" usually attaches to a log that is
-				// already closed and the stream ends before anything arrives.
-				// Waiting first makes the output deterministic.
-				if err := client.StreamBatchExecutionLogs(cfg.Project(), cfg.Service(), exec.ID, os.Stdout, false); err != nil {
+				if err := followBatchOutput(client, cfg, exec.ID); err != nil {
 					fmt.Fprintf(os.Stderr, "could not read execution output: %v\n", err)
 				}
 			}
+			final, err := awaitBatchCompletion(client, cfg, exec.ID)
+			if err != nil {
+				return err
+			}
 			fmt.Printf("Execution %s %s (exit %d)\n", final.ID, colorize(final.Status), final.ExitCode)
 			if final.Status != "succeeded" {
 				if final.Error != "" {
@@ -160,7 +155,11 @@ latest successful build is used.`,
 			if err != nil {
 				return err
 			}
-			return client.StreamBatchExecutionLogs(cfg.Project(), cfg.Service(), args[0], os.Stdout, batchLogsFollow)
+			if batchLogsFollow {
+				return followBatchOutput(client, cfg, args[0])
+			}
+			_, err = client.ReadBatchExecutionLogs(cfg.Project(), cfg.Service(), args[0], 0, os.Stdout)
+			return err
 		}),
 	}
 
@@ -299,28 +298,83 @@ func shortID(id string) string {
 	return id
 }
 
-const batchPollInterval = 2 * time.Second
+const (
+	batchPollInterval = 2 * time.Second
+	// Log tailing polls faster than status, so output feels live.
+	batchLogPollInterval = 750 * time.Millisecond
+)
 
 func awaitBatchCompletion(client *api.Client, cfg config.ServiceConfig, id string) (*api.BatchExecution, error) {
 	for {
 		exec, err := client.InspectBatchExecution(cfg.Project(), cfg.Service(), id)
 		if err != nil {
 			return nil, err
 		}
-		switch exec.Status {
-		case "succeeded", "failed", "cancelled":
+		if isTerminalBatchStatus(exec.Status) {
 			return exec, nil
 		}
 		time.Sleep(batchPollInterval)
 	}
 }
 
+// followBatchOutput shows an execution's output as it is produced, returning
+// once the execution has finished and its output has been fully drained.
+//
+// It polls from a byte offset rather than holding a server-side follow open. The
+// server never closes a followed batch log, so a follower cannot tell when the
+// execution is done and simply hangs -- which it did. Polling puts the
+// termination condition where the answer actually lives: the execution's status.
+func followBatchOutput(client *api.Client, cfg config.ServiceConfig, id string) error {
+	offset := 0
+	for {
+		exec, err := client.InspectBatchExecution(cfg.Project(), cfg.Service(), id)
+		if err != nil {
+			return err
+		}
+		// A queued execution has no log yet, and asking for one is an error
+		// rather than an empty read.
+		if exec.Status != "pending" {
+			n, err := client.ReadBatchExecutionLogs(cfg.Project(), cfg.Service(), id, offset, os.Stdout)
+			offset += n
+			if err != nil && !isTerminalBatchStatus(exec.Status) {
+				// Transient while the container is still coming up.
+				time.Sleep(batchPollInterval)
+				continue
+			} else if err != nil {
+				return err
+			}
+		}
+		if isTerminalBatchStatus(exec.Status) {
+			// The status was read before the log, so anything written between
+			// the two reads is still outstanding.
+			n, err := client.ReadBatchExecutionLogs(cfg.Project(), cfg.Service(), id, offset, os.Stdout)
+			offset += n
+			if err != nil {
+				return err
+			}
+			if n == 0 {
+				return nil
+			}
+			continue
+		}
+		time.Sleep(batchLogPollInterval)
+	}
+}
+
+func isTerminalBatchStatus(status string) bool {
+	switch status {
+	case "succeeded", "failed", "cancelled":
+		return true
+	}
+	return false
+}
+
 func initBatchCmd() {
 	batchCmd.PersistentFlags().StringVarP(&batchService, "service", "s", "", "The batch service to operate on")
 
 	batchRunCmd.Flags().StringVarP(&batchRunBuild, "build", "b", "", "Build to run (default: latest successful)")
 	batchRunCmd.Flags().StringArrayVarP(&batchRunEnv, "env", "e", nil, "Environment variable as KEY=VALUE, repeatable")
-	batchRunCmd.Flags().BoolVarP(&batchRunFollow, "follow", "f", false, "Wait for the execution to finish, then print its output")
+	batchRunCmd.Flags().BoolVarP(&batchRunFollow, "follow", "f", false, "Show the output as the execution runs, and wait for it to finish")
 	batchRunCmd.Flags().BoolVarP(&batchRunWait, "wait", "w", false, "Wait for the execution to finish and exit non-zero if it failed")
 	batchLogsCmd.Flags().BoolVarP(&batchLogsFollow, "follow", "f", false, "Follow the output")