diff --git a/cmd/src/batch_common.go b/cmd/src/batch_common.go index 6a78d8f296..1fd867f8c2 100644 --- a/cmd/src/batch_common.go +++ b/cmd/src/batch_common.go @@ -12,6 +12,7 @@ import ( "path/filepath" "runtime" "strings" + "syscall" "time" "github.com/mattn/go-isatty" @@ -220,7 +221,7 @@ func batchDefaultTempDirPrefix() string { return os.TempDir() } -func batchOpenFileFlag(flag string) (io.ReadCloser, error) { +func batchOpenFileFlag(flag string) (*os.File, error) { if flag == "" || flag == "-" { if flag != "-" { // If the flag wasn't set, we want to check stdin. If it's not a TTY, @@ -238,8 +239,13 @@ func batchOpenFileFlag(flag string) (io.ReadCloser, error) { } } } + // https://github.com/golang/go/issues/24842 + if err := syscall.SetNonblock(0, true); err != nil { + panic(err) + } + stdin := os.NewFile(0, "stdin") - return os.Stdin, nil + return stdin, nil } file, err := os.Open(flag) @@ -292,7 +298,7 @@ func executeBatchSpec(ctx context.Context, ui ui.ExecUI, opts executeBatchSpecOp // Parse flags and build up our service and executor options. ui.ParsingBatchSpec() - batchSpec, rawSpec, err := parseBatchSpec(opts.file, svc, false) + batchSpec, rawSpec, err := parseBatchSpec(ctx, opts.file, svc, false) if err != nil { var multiErr errors.MultiError if errors.As(err, &multiErr) { @@ -481,18 +487,33 @@ func executeBatchSpec(ctx context.Context, ui ui.ExecUI, opts executeBatchSpecOp return nil } +type ReadDeadliner interface { + SetReadDeadline(t time.Time) error +} + +func SetReadDeadlineOnCancel(ctx context.Context, d ReadDeadliner) { + go func() { + <-ctx.Done() + d.SetReadDeadline(time.Now()) + }() +} + // parseBatchSpec parses and validates the given batch spec. If the spec has // validation errors, they are returned. // // isRemote argument is a temporary argument used to determine if the batch spec is being parsed for remote // (server-side) processing. Remote processing does not support mounts yet. -func parseBatchSpec(file string, svc *service.Service, isRemote bool) (*batcheslib.BatchSpec, string, error) { +func parseBatchSpec(ctx context.Context, file string, svc *service.Service, isRemote bool) (*batcheslib.BatchSpec, string, error) { f, err := batchOpenFileFlag(file) if err != nil { return nil, "", err } defer f.Close() + ctx, cancel := context.WithCancel(ctx) + defer cancel() + SetReadDeadlineOnCancel(ctx, f) + data, err := io.ReadAll(f) if err != nil { return nil, "", errors.Wrap(err, "reading batch spec") diff --git a/cmd/src/batch_remote.go b/cmd/src/batch_remote.go index 9d542ae9b5..261831c11c 100644 --- a/cmd/src/batch_remote.go +++ b/cmd/src/batch_remote.go @@ -63,7 +63,7 @@ Examples: // may as well validate it at the same time so we don't even have to go to // the backend if it's invalid. ui.ParsingBatchSpec() - spec, raw, err := parseBatchSpec(file, svc, true) + spec, raw, err := parseBatchSpec(ctx, file, svc, true) if err != nil { ui.ParsingBatchSpecFailure(err) return err diff --git a/cmd/src/batch_repositories.go b/cmd/src/batch_repositories.go index 83beb40915..8adbfd9643 100644 --- a/cmd/src/batch_repositories.go +++ b/cmd/src/batch_repositories.go @@ -75,7 +75,7 @@ Examples: } out := output.NewOutput(flagSet.Output(), output.OutputOpts{Verbose: *verbose}) - spec, _, err := parseBatchSpec(file, svc, false) + spec, _, err := parseBatchSpec(ctx, file, svc, false) if err != nil { ui := &ui.TUI{Out: out} ui.ParsingBatchSpecFailure(err) diff --git a/cmd/src/batch_validate.go b/cmd/src/batch_validate.go index 3d40a8caca..db0e9d73e4 100644 --- a/cmd/src/batch_validate.go +++ b/cmd/src/batch_validate.go @@ -75,7 +75,7 @@ Examples: return err } - if _, _, err := parseBatchSpec(file, svc, false); err != nil { + if _, _, err := parseBatchSpec(ctx, file, svc, false); err != nil { ui.ParsingBatchSpecFailure(err) return err }