diff --git a/README.md b/README.md index 038e31e..7e6e332 100644 --- a/README.md +++ b/README.md @@ -322,6 +322,7 @@ These are methods on a pipe that change its configuration: | Source | Modifies | | -------- | ------------- | +| [`WithContext`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithContext) | context used in commands or HTTP requests | | [`WithEnv`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithEnv) | environment for commands | | [`WithError`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithError) | pipe error status | | [`WithHTTPClient`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithHTTPClient) | client for HTTP requests | diff --git a/script.go b/script.go index 2195da4..457505d 100644 --- a/script.go +++ b/script.go @@ -3,6 +3,7 @@ package script import ( "bufio" "container/ring" + "context" "crypto/sha256" "encoding/base64" "encoding/hex" @@ -38,6 +39,7 @@ type Pipe struct { err error stderr io.Writer env []string + ctx context.Context } // Args creates a pipe containing the program's command-line arguments from @@ -174,6 +176,7 @@ func NewPipe() *Pipe { stdout: os.Stdout, httpClient: http.DefaultClient, env: nil, + ctx: context.Background(), } } @@ -198,6 +201,11 @@ func Stdin() *Pipe { return NewPipe().WithReader(os.Stdin) } +// WithContext creates a pipe with the context ctx. +func WithContext(ctx context.Context) *Pipe { + return NewPipe().WithContext(ctx) +} + // AppendFile appends the contents of the pipe to the file path, creating it if // necessary, and returns the number of bytes successfully written, or an // error. @@ -326,6 +334,12 @@ func (p *Pipe) Dirname() *Pipe { // set by [Pipe.WithHTTPClient], or [http.DefaultClient] otherwise. The // response body is streamed concurrently to the pipe's output. If the response // status is anything other than HTTP 200-299, the pipe's error status is set. +// +// # Context +// +// The request does not inherit the pipe's context (if any was set by +// [Pipe.WithContext]). If you want the request to be cancellable, you will need to set +// your own context on it. func (p *Pipe) Do(req *http.Request) *Pipe { return p.Filter(func(r io.Reader, w io.Writer) error { resp, err := p.httpClient.Do(req) @@ -413,6 +427,11 @@ func (p *Pipe) Error() error { // The command inherits the current process's environment, optionally modified // by [Pipe.WithEnv]. // +// # Context +// +// The command inherits the pipe's context (if any was set by [Pipe.WithContext]), and +// will be cancelled if the context is cancelled or times out. +// // # Error handling // // If the command had a non-zero exit status, the pipe's error status will also @@ -432,7 +451,7 @@ func (p *Pipe) Exec(cmdLine string) *Pipe { if err != nil { return err } - cmd := exec.Command(args[0], args[1:]...) + cmd := exec.CommandContext(p.ctx, args[0], args[1:]...) cmd.Stdin = r cmd.Stdout = w cmd.Stderr = w @@ -462,6 +481,16 @@ func (p *Pipe) Exec(cmdLine string) *Pipe { // syntax. For example: // // ListFiles("*").ExecForEach("touch {{.}}").Wait() +// +// # Environment +// +// The command inherits the current process's environment, optionally modified +// by [Pipe.WithEnv]. +// +// # Context +// +// The command inherits the pipe's context (if any was set by [Pipe.WithContext]), and +// will be cancelled if the context is cancelled or times out. func (p *Pipe) ExecForEach(cmdLine string) *Pipe { tpl, err := template.New("").Parse(cmdLine) if err != nil { @@ -470,6 +499,9 @@ func (p *Pipe) ExecForEach(cmdLine string) *Pipe { return p.Filter(func(r io.Reader, w io.Writer) error { scanner := newScanner(r) for scanner.Scan() { + if p.ctx.Err() != nil { + return p.ctx.Err() + } cmdLine := new(strings.Builder) err := tpl.Execute(cmdLine, scanner.Text()) if err != nil { @@ -479,7 +511,7 @@ func (p *Pipe) ExecForEach(cmdLine string) *Pipe { if err != nil { return err } - cmd := exec.Command(args[0], args[1:]...) + cmd := exec.CommandContext(p.ctx, args[0], args[1:]...) cmd.Stdout = w cmd.Stderr = w pipeStderr := p.stdErr() @@ -530,9 +562,10 @@ func (p *Pipe) ExitStatus() int { // [io.Writer] to write its output to, and returns an error, which will be set // on the pipe. // -// filter runs concurrently, so its goroutine will not exit until the pipe has -// been fully read. Use [Pipe.Wait] to wait for all concurrent filters to -// complete. +// filter runs concurrently, so its goroutine will not exit until the pipe has been +// fully read. Use [Pipe.Wait] to wait for all concurrent filters to complete. These +// goroutines are not automatically cancelled by the pipe's context, if one was set by +// [Pipe.WithContext]. func (p *Pipe) Filter(filter func(io.Reader, io.Writer) error) *Pipe { if p.Error() != nil { return p @@ -651,8 +684,13 @@ func (p *Pipe) Freq() *Pipe { // Get makes an HTTP GET request to url, sending the contents of the pipe as // the request body, and produces the server's response. See [Pipe.Do] for how // the HTTP response status is interpreted. +// +// # Context +// +// The request inherits the pipe's context (if any was set by [Pipe.WithContext]), and +// will be cancelled if the context is cancelled or times out. func (p *Pipe) Get(url string) *Pipe { - req, err := http.NewRequest(http.MethodGet, url, p.Reader) + req, err := http.NewRequestWithContext(p.ctx, http.MethodGet, url, p.Reader) if err != nil { return p.WithError(err) } @@ -806,8 +844,13 @@ func (p *Pipe) MatchRegexp(re *regexp.Regexp) *Pipe { // Post makes an HTTP POST request to url, using the contents of the pipe as // the request body, and produces the server's response. See [Pipe.Do] for how // the HTTP response status is interpreted. +// +// # Context +// +// The request inherits the pipe's context (if any was set by [Pipe.WithContext]), and +// will be cancelled if the context is cancelled or times out. func (p *Pipe) Post(url string) *Pipe { - req, err := http.NewRequest(http.MethodPost, url, p.Reader) + req, err := http.NewRequestWithContext(p.ctx, http.MethodPost, url, p.Reader) if err != nil { return p.WithError(err) } @@ -963,6 +1006,12 @@ func (p *Pipe) Wait() error { return p.Error() } +// WithContext sets the context for subsequent [Pipe.Exec], [Pipe.ExecForEach], [Pipe.Get] and [Pipe.Post] commands. +func (p *Pipe) WithContext(ctx context.Context) *Pipe { + p.ctx = ctx + return p +} + // WithEnv sets the environment for subsequent [Pipe.Exec] and [Pipe.ExecForEach] // commands to the string slice env, using the same format as [os/exec.Cmd.Env]. // An empty slice unsets all existing environment variables. diff --git a/script_test.go b/script_test.go index 66e4df5..c4f9a77 100644 --- a/script_test.go +++ b/script_test.go @@ -3,6 +3,7 @@ package script_test import ( "bufio" "bytes" + "context" "crypto/sha256" "crypto/sha512" "errors" @@ -18,6 +19,7 @@ import ( "strings" "testing" "testing/iotest" + "time" "github.com/bitfield/script" "github.com/google/go-cmp/cmp" @@ -2165,6 +2167,7 @@ func TestHash_ReturnsErrorGivenReadErrorOnPipe(t *testing.T) { } func TestHashSums_OutputsEmptyStringForFileThatCannotBeHashed(t *testing.T) { + t.Parallel() got, err := script.Echo("file_does_not_exist.txt").HashSums(sha256.New()).String() if err != nil { t.Fatal(err) @@ -2175,6 +2178,55 @@ func TestHashSums_OutputsEmptyStringForFileThatCannotBeHashed(t *testing.T) { } } +func TestWithContext_SkipsExecWithCancelledContext(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + cancel() + output, err := script.WithContext(ctx).Exec("echo Hello, World").String() + if err == nil { + t.Error("expected error due to context cancellation, got nil") + } + if strings.Contains(output, "Hello") { + t.Fatal("command ran even though context cancelled beforehand") + } +} + +func TestWithContext_SkipsGETRequestWithCancelledContext(t *testing.T) { + t.Parallel() + handler_called := false + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + handler_called = true + })) + defer ts.Close() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := script.WithContext(ctx).Echo("request data").Get(ts.URL).String() + if err == nil { + t.Fatal("expected error from cancelled context") + } + if handler_called { + t.Fatal("request made even though context cancelled beforehand") + } +} + +func TestWithContext_SkipsPOSTRequestWithCancelledContext(t *testing.T) { + t.Parallel() + handler_called := false + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + handler_called = true + })) + defer ts.Close() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := script.WithContext(ctx).Echo("request data").Post(ts.URL).String() + if err == nil { + t.Fatal("expected error from cancelled context") + } + if handler_called { + t.Fatal("request made even though context cancelled beforehand") + } +} + func ExampleArgs() { script.Args().Stdout() // prints command-line arguments @@ -2622,6 +2674,21 @@ func ExamplePipe_Tee_writers() { // hello } +func ExamplePipe_WithContext() { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(20 * time.Millisecond) + })) + defer ts.Close() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + _, err := script.NewPipe().WithContext(ctx).Get(ts.URL).String() + if err != nil { + fmt.Println("cancelled") + } + // Output: + // cancelled +} + func ExamplePipe_WithStderr() { buf := new(bytes.Buffer) script.NewPipe().WithStderr(buf).Exec("go").Wait()