Skip to content

Commit

Permalink
Cancel reading from stdin on Ctrl-C
Browse files Browse the repository at this point in the history
This is essentially a copy-paste implementation of the ideas presented
in this comment here: golang/go#20280 (comment)

It fixes #775 and helps with the issue described in #793 (comment). Not sure if it has unintended side-effects.
  • Loading branch information
mrnugget authored and BolajiOlajide committed Jul 7, 2022
1 parent 05e3f02 commit 219d25e
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 7 deletions.
29 changes: 25 additions & 4 deletions cmd/src/batch_common.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (
"path/filepath"
"runtime"
"strings"
"syscall"
"time"

"github.com/mattn/go-isatty"
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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")
Expand Down
2 changes: 1 addition & 1 deletion cmd/src/batch_remote.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion cmd/src/batch_repositories.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion cmd/src/batch_validate.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down

0 comments on commit 219d25e

Please sign in to comment.