IOConnectorPair must close its connectors at the end

Before this change, the connectors were closed by the cancelletion
of its associated (passed) context, which seems too early.

Signed-off-by: Kazuyoshi Kato <katokazu@amazon.com>

Author: Kazuyoshi Kato

internal/vm/ioproxy.go

@@ -67,20 +67,23 @@ type IOConnectorPair struct {
 }
 
 func (connectorPair *IOConnectorPair) proxy(
-	proc *vmProc,
+	ctx context.Context,
 	logger *logrus.Entry,
 	timeoutAfterExit time.Duration,
 ) (ioInitDone <-chan error, ioCopyDone <-chan error) {
 	initDone := make(chan error, 2)
 	copyDone := make(chan error)
 
+	ioCtx, ioCancel := context.WithCancel(context.Background())
+
 	// Start the initialization process. Any synchronous setup made by the connectors will
 	// be completed after these lines. Async setup will be done once initDone is closed in
 	// the goroutine below.
-	readerResultCh := connectorPair.ReadConnector(proc.ctx, logger.WithField("direction", "read"))
-	writerResultCh := connectorPair.WriteConnector(proc.ctx, logger.WithField("direction", "write"))
+	readerResultCh := connectorPair.ReadConnector(ioCtx, logger.WithField("direction", "read"))
+	writerResultCh := connectorPair.WriteConnector(ioCtx, logger.WithField("direction", "write"))
 
 	go func() {
+		defer ioCancel()
 		defer close(copyDone)
 
 		var reader io.ReadCloser
@@ -119,7 +122,7 @@ func (connectorPair *IOConnectorPair) proxy(
 		// If the io streams close on their own before the timeout, the Close calls here
 		// should just be no-ops.
 		go func() {
-			<-proc.ctx.Done()
+			<-ctx.Done()
 			time.AfterFunc(timeoutAfterExit, func() {
 				logClose(logger, reader, writer)
 			})
@@ -129,6 +132,7 @@ func (connectorPair *IOConnectorPair) proxy(
 		defer logger.Debug("end copying io")
 
 		size, err := io.CopyBuffer(writer, reader, make([]byte, internal.DefaultBufferSize))
+		logger.Debugf("copied %d", size)
 		if err != nil {
 			if strings.Contains(err.Error(), "use of closed network connection") ||
 				strings.Contains(err.Error(), "file already closed") {
@@ -138,7 +142,6 @@ func (connectorPair *IOConnectorPair) proxy(
 			}
 			copyDone <- err
 		}
-		logger.Debugf("copied %d", size)
 		defer logClose(logger, reader, writer)
 	}()
 
@@ -174,19 +177,20 @@ func (ioConnectorSet *ioConnectorSet) start(proc *vmProc) (ioInitDone <-chan err
 	if ioConnectorSet.stdin != nil {
 		// For Stdin only, provide 0 as the timeout to wait after the proc exits before closing IO streams.
 		// There's no reason to send stdin data to a proc that's already dead.
-		waitErrs(ioConnectorSet.stdin.proxy(proc, proc.logger.WithField("stream", "stdin"), 0))
+		waitErrs(ioConnectorSet.stdin.proxy(proc.ctx, proc.logger.WithField("stream", "stdin"), 0))
+
 	} else {
 		proc.logger.Debug("skipping proxy io for unset stdin")
 	}
 
 	if ioConnectorSet.stdout != nil {
-		waitErrs(ioConnectorSet.stdout.proxy(proc, proc.logger.WithField("stream", "stdout"), defaultIOFlushTimeout))
+		waitErrs(ioConnectorSet.stdout.proxy(proc.ctx, proc.logger.WithField("stream", "stdout"), defaultIOFlushTimeout))
 	} else {
 		proc.logger.Debug("skipping proxy io for unset stdout")
 	}
 
 	if ioConnectorSet.stderr != nil {
-		waitErrs(ioConnectorSet.stderr.proxy(proc, proc.logger.WithField("stream", "stderr"), defaultIOFlushTimeout))
+		waitErrs(ioConnectorSet.stderr.proxy(proc.ctx, proc.logger.WithField("stream", "stderr"), defaultIOFlushTimeout))
 	} else {
 		proc.logger.Debug("skipping proxy io for unset stderr")
 	}

internal/vm/ioproxy_test.go

@@ -0,0 +1,66 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package vm
+
+import (
+	"context"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func fileConnector(path string, flag int) IOConnector {
+	return func(procCtx context.Context, logger *logrus.Entry) <-chan IOConnectorResult {
+		returnCh := make(chan IOConnectorResult, 1)
+		defer close(returnCh)
+
+		file, err := os.OpenFile(path, flag, 0600)
+		returnCh <- IOConnectorResult{
+			ReadWriteCloser: file,
+			Err:             err,
+		}
+
+		return returnCh
+	}
+}
+
+func TestProxy(t *testing.T) {
+	dir, err := ioutil.TempDir("", t.Name())
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+
+	ctx := context.Background()
+	content := "hello world"
+
+	err = ioutil.WriteFile(filepath.Join(dir, "input"), []byte(content), 0600)
+	require.NoError(t, err)
+
+	pair := &IOConnectorPair{
+		ReadConnector:  fileConnector(filepath.Join(dir, "input"), os.O_RDONLY),
+		WriteConnector: fileConnector(filepath.Join(dir, "output"), os.O_CREATE|os.O_WRONLY),
+	}
+	initCh, copyCh := pair.proxy(ctx, logrus.WithFields(logrus.Fields{}), 0)
+
+	assert.Nil(t, <-initCh)
+	assert.Nil(t, <-copyCh)
+
+	bytes, err := ioutil.ReadFile(filepath.Join(dir, "output"))
+	require.NoError(t, err)
+	assert.Equal(t, content, string(bytes))
+}

internal/vm/task.go

@@ -269,9 +269,6 @@ func (m *taskManager) monitorExit(proc *vmProc, taskService taskAPI.TaskService)
 		ID:     proc.taskID,
 		ExecID: proc.execID,
 	})
-
-	<-proc.ioCopyDone
-
 	proc.cancel()
 
 	if waitErr == context.Canceled {

runtime/service_integ_test.go

@@ -1589,7 +1589,7 @@ func TestExec_Isolated(t *testing.T) {
 	var taskStdout bytes.Buffer
 	var taskStderr bytes.Buffer
 
-	task, err := c.NewTask(ctx, cio.NewCreator(cio.WithStreams(nil, &taskStdout, &taskStderr)))
+	task, err := c.NewTask(ctx, cio.NewCreator(cio.WithStreams(os.Stdin, &taskStdout, &taskStderr)))
 	require.NoError(t, err, "failed to create task for container %s", c.ID())
 
 	taskExitCh, err := task.Wait(ctx)
@@ -1604,7 +1604,7 @@ func TestExec_Isolated(t *testing.T) {
 	taskExec, err := task.Exec(ctx, "exec", &specs.Process{
 		Args: []string{"/bin/date"},
 		Cwd:  "/",
-	}, cio.NewCreator(cio.WithStreams(nil, &execStdout, &execStderr)))
+	}, cio.NewCreator(cio.WithStreams(os.Stdin, &execStdout, &execStderr)))
 	require.NoError(t, err)
 
 	execExitCh, err := taskExec.Wait(ctx)