Skip to content

Instantly share code, notes, and snippets.

@dreampuf
Created November 8, 2018 22:15
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save dreampuf/470a07e41211691e4e69349ae45c991b to your computer and use it in GitHub Desktop.
Save dreampuf/470a07e41211691e4e69349ae45c991b to your computer and use it in GitHub Desktop.
How to handle sub process gently in Golang
package main
import (
"bufio"
"bytes"
"context"
"io"
"log"
"os/exec"
"strings"
"syscall"
"testing"
"time"
)
func Test_OsExecBasic(t *testing.T) {
cmd := exec.Command("sleep", ".1")
if err := cmd.Run(); err != nil {
t.Error(err)
}
cmdCtx, _ := context.WithTimeout(context.TODO(), time.Millisecond * 200)
cmdWithCtx := exec.CommandContext(cmdCtx, "sleep", ".3")
if err := cmdWithCtx.Run(); err != nil && err.Error() != "signal: killed" {
t.Error(err)
}
}
func Test_OsExecInputAndOutput(t *testing.T) {
cmdCtx, _ := context.WithCancel(context.TODO())
outputBuf := bytes.NewBuffer(nil)
case01 := "I'm input"
cmd := exec.CommandContext(cmdCtx, "cat", "-")
cmd.Stdin = strings.NewReader(case01)
cmd.Stdout = outputBuf
if err := cmd.Run(); err != nil {
t.Error(err)
}
if outputBuf.String() != case01 {
t.Error("output isn't equal to input")
}
}
func Test_OsExecEnvironment(t *testing.T) {
cmdCtx, _ := context.WithCancel(context.TODO())
// You can't use echo directly in Mac OSX
// https://apple.stackexchange.com/questions/173836/why-echo-n-doesnt-work-in-this-script-on-mac-terminal
cmd := exec.CommandContext(cmdCtx, "sh", "-c", "/bin/echo -n $FOO")
cmd.Env = []string{"FOO=BAR"}
output, err := cmd.CombinedOutput()
if err != nil {
t.Error(err)
}
if string(output) != "BAR" {
t.Errorf("'%s' != 'BAR'", string(output))
}
}
func holdUntilStart(cmd *exec.Cmd) {
for {
// in case output not ready
if cmd.ProcessState != nil {
break
}
time.Sleep(time.Millisecond)
}
}
func Test_HandleLargeOutput(t *testing.T) {
cmdCtx, _ := context.WithCancel(context.TODO())
case01 := "a huge string\na huge string\na huge string"
outputBuf := bytes.NewBuffer(nil)
onHandleFinished := make(chan struct{})
cmd := exec.CommandContext(cmdCtx, "cat", "-")
cmd.Stdin = strings.NewReader(case01)
cmd.Stdout = outputBuf
if err := cmd.Start(); err != nil {
t.Error(err)
}
go func() {
holdUntilStart(cmd)
scanner := bufio.NewScanner(outputBuf)
for scanner.Scan() {
// hanlde scanner.Text()
//log.Print(scanner.Text())
}
if scanner.Err() != nil {
t.Error(scanner.Err())
}
close(onHandleFinished)
}()
if err := cmd.Wait(); err != nil {
t.Error(err)
}
<- onHandleFinished
}
func Test_HandleProcAttr(t *testing.T) {
cmdCtx, _ := context.WithTimeout(context.TODO(), time.Second)
cmd := exec.CommandContext(cmdCtx, "sh", "-c", "sleep 5 &\nwait")
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
if err := cmd.Start(); err != nil {
t.Error(err)
}
time.AfterFunc(time.Second, func() {
// Kill all sub processes
// https://medium.com/@felixge/killing-a-child-process-and-all-of-its-children-in-go-54079af94773
syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
})
if err := cmd.Wait(); err != nil && err.Error() != "signal: killed" {
t.Error(err)
}
}
func Test_StdoutStreamingProcess(t *testing.T) {
cmdCtx, _ := context.WithTimeout(context.TODO(), time.Second)
pr, pw := io.Pipe()
cmd := exec.CommandContext(cmdCtx, "cat", "-")
cmd.Stdin = pr
cmd.Stdout = pw
onHandleFinished := make(chan struct{})
if err := cmd.Start(); err != nil {
t.Error(err)
}
go func() {
count := 0
for {
if n, err := pw.Write([]byte{'a', 'b', 'c'}); err != nil {
if err != io.ErrClosedPipe {
t.Error(err)
}
break
} else {
count += n
}
}
log.Printf("Write %d bytes\n", count)
}()
go func() {
scanner := bufio.NewScanner(pr)
scanner.Split(bufio.ScanBytes)
for scanner.Scan() {
//log.Print(scanner.Text())
}
pr.Close()
close(onHandleFinished)
}()
go func() {
if err := cmd.Wait(); err != nil && err.Error() != "signal: killed" {
t.Error(err)
}
log.Println("wait done")
pw.Close()
}()
<- onHandleFinished
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment