From 814fec60f08d1992d8d049f943a9c5c2e383528e Mon Sep 17 00:00:00 2001
From: Christian Schwarz <me@cschwarz.com>
Date: Fri, 19 Oct 2018 16:12:21 +0200
Subject: [PATCH] endpoint + zfs: context cancellation of util.IOCommand
 instances (send & recv for now)

---
 endpoint/endpoint.go |  4 ++--
 util/iocommand.go    |  9 +++++----
 zfs/zfs.go           | 29 ++++-------------------------
 3 files changed, 11 insertions(+), 31 deletions(-)

diff --git a/endpoint/endpoint.go b/endpoint/endpoint.go
index 9202fa9..be3a00b 100644
--- a/endpoint/endpoint.go
+++ b/endpoint/endpoint.go
@@ -89,7 +89,7 @@ func (p *Sender) Send(ctx context.Context, r *pdu.SendReq) (*pdu.SendRes, io.Rea
 		}
 		return &pdu.SendRes{ExpectedSize: expSize}, nil, nil
 	} else {
-		stream, err := zfs.ZFSSend(r.Filesystem, r.From, r.To, "")
+		stream, err := zfs.ZFSSend(ctx, r.Filesystem, r.From, r.To, "")
 		if err != nil {
 			return nil, nil, err
 		}
@@ -279,7 +279,7 @@ func (e *Receiver) Receive(ctx context.Context, req *pdu.ReceiveReq, sendStream
 
 	getLogger(ctx).Debug("start receive command")
 
-	if err := zfs.ZFSRecv(lp.ToString(), sendStream, args...); err != nil {
+	if err := zfs.ZFSRecv(ctx, lp.ToString(), sendStream, args...); err != nil {
 		getLogger(ctx).
 			WithError(err).
 			WithField("args", args).
diff --git a/util/iocommand.go b/util/iocommand.go
index 0433446..1113f74 100644
--- a/util/iocommand.go
+++ b/util/iocommand.go
@@ -2,6 +2,7 @@ package util
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -34,8 +35,8 @@ func (e IOCommandError) Error() string {
 	return fmt.Sprintf("underlying process exited with error: %s\nstderr: %s\n", e.WaitErr, e.Stderr)
 }
 
-func RunIOCommand(command string, args ...string) (c *IOCommand, err error) {
-	c, err = NewIOCommand(command, args, IOCommandStderrBufSize)
+func RunIOCommand(ctx context.Context, command string, args ...string) (c *IOCommand, err error) {
+	c, err = NewIOCommand(ctx, command, args, IOCommandStderrBufSize)
 	if err != nil {
 		return
 	}
@@ -43,7 +44,7 @@ func RunIOCommand(command string, args ...string) (c *IOCommand, err error) {
 	return
 }
 
-func NewIOCommand(command string, args []string, stderrBufSize int) (c *IOCommand, err error) {
+func NewIOCommand(ctx context.Context, command string, args []string, stderrBufSize int) (c *IOCommand, err error) {
 
 	if stderrBufSize == 0 {
 		stderrBufSize = IOCommandStderrBufSize
@@ -51,7 +52,7 @@ func NewIOCommand(command string, args []string, stderrBufSize int) (c *IOComman
 
 	c = &IOCommand{}
 
-	c.Cmd = exec.Command(command, args...)
+	c.Cmd = exec.CommandContext(ctx, command, args...)
 
 	if c.Stdout, err = c.Cmd.StdoutPipe(); err != nil {
 		return
diff --git a/zfs/zfs.go b/zfs/zfs.go
index 3890449..40a0ea8 100644
--- a/zfs/zfs.go
+++ b/zfs/zfs.go
@@ -317,7 +317,7 @@ func buildCommonSendArgs(fs string, from, to string, token string) ([]string, er
 // if token != "", then send -t token is used
 // otherwise send [-i from] to is used
 // (if from is "" a full ZFS send is done)
-func ZFSSend(fs string, from, to string, token string) (stream io.ReadCloser, err error) {
+func ZFSSend(ctx context.Context, fs string, from, to string, token string) (stream io.ReadCloser, err error) {
 
 	args := make([]string, 0)
 	args = append(args, "send")
@@ -328,7 +328,7 @@ func ZFSSend(fs string, from, to string, token string) (stream io.ReadCloser, er
 	}
 	args = append(args, sargs...)
 
-	stream, err = util.RunIOCommand(ZFS_BINARY, args...)
+	stream, err = util.RunIOCommand(ctx, ZFS_BINARY, args...)
 
 	return
 }
@@ -455,7 +455,7 @@ func ZFSSendDry(fs string, from, to string, token string) (_ *DrySendInfo, err e
 }
 
 
-func ZFSRecv(fs string, stream io.Reader, additionalArgs ...string) (err error) {
+func ZFSRecv(ctx context.Context, fs string, stream io.Reader, additionalArgs ...string) (err error) {
 
 	if err := validateZFSFilesystem(fs); err != nil {
 		return err
@@ -468,7 +468,7 @@ func ZFSRecv(fs string, stream io.Reader, additionalArgs ...string) (err error)
 	}
 	args = append(args, fs)
 
-	cmd := exec.Command(ZFS_BINARY, args...)
+	cmd := exec.CommandContext(ctx, ZFS_BINARY, args...)
 
 	stderr := bytes.NewBuffer(make([]byte, 0, 1024))
 	cmd.Stderr = stderr
@@ -523,27 +523,6 @@ func ZFSRecvClearResumeToken(fs string) (err error) {
 	return nil
 }
 
-func ZFSRecvWriter(fs *DatasetPath, additionalArgs ...string) (io.WriteCloser, error) {
-
-	args := make([]string, 0)
-	args = append(args, "recv")
-	if len(args) > 0 {
-		args = append(args, additionalArgs...)
-	}
-	args = append(args, fs.ToString())
-
-	cmd, err := util.NewIOCommand(ZFS_BINARY, args, 1024)
-	if err != nil {
-		return nil, err
-	}
-
-	if err = cmd.Start(); err != nil {
-		return nil, err
-	}
-
-	return cmd.Stdin, nil
-}
-
 type ZFSProperties struct {
 	m map[string]string
 }