Merge pull request #421 from kzys/exec-test-master

Fix Exec()'s timeout issue

Author: Kazuyoshi Kato

agent/service.go

@@ -228,15 +228,15 @@ func (ts *TaskService) Create(requestCtx context.Context, req *taskAPI.CreateTas
 			req.Stdin = fifoSet.Stdin
 			stdinConnectorPair = &vm.IOConnectorPair{
 				ReadConnector:  vm.VSockAcceptConnector(extraData.StdinPort),
-				WriteConnector: vm.FIFOConnector(fifoSet.Stdin),
+				WriteConnector: vm.WriteFIFOConnector(fifoSet.Stdin),
 			}
 		}
 
 		var stdoutConnectorPair *vm.IOConnectorPair
 		if req.Stdout != "" {
 			req.Stdout = fifoSet.Stdout
 			stdoutConnectorPair = &vm.IOConnectorPair{
-				ReadConnector:  vm.FIFOConnector(fifoSet.Stdout),
+				ReadConnector:  vm.ReadFIFOConnector(fifoSet.Stdout),
 				WriteConnector: vm.VSockAcceptConnector(extraData.StdoutPort),
 			}
 		}
@@ -245,7 +245,7 @@ func (ts *TaskService) Create(requestCtx context.Context, req *taskAPI.CreateTas
 		if req.Stderr != "" {
 			req.Stderr = fifoSet.Stderr
 			stderrConnectorPair = &vm.IOConnectorPair{
-				ReadConnector:  vm.FIFOConnector(fifoSet.Stderr),
+				ReadConnector:  vm.ReadFIFOConnector(fifoSet.Stderr),
 				WriteConnector: vm.VSockAcceptConnector(extraData.StderrPort),
 			}
 		}
@@ -454,7 +454,7 @@ func (ts *TaskService) Exec(requestCtx context.Context, req *taskAPI.ExecProcess
 			req.Stdin = fifoSet.Stdin
 			stdinConnectorPair = &vm.IOConnectorPair{
 				ReadConnector:  vm.VSockAcceptConnector(extraData.StdinPort),
-				WriteConnector: vm.FIFOConnector(fifoSet.Stdin),
+				WriteConnector: vm.WriteFIFOConnector(fifoSet.Stdin),
 			}
 			ts.addCleanup(taskExecID, func() error {
 				return os.RemoveAll(req.Stdin)
@@ -465,7 +465,7 @@ func (ts *TaskService) Exec(requestCtx context.Context, req *taskAPI.ExecProcess
 		if req.Stdout != "" {
 			req.Stdout = fifoSet.Stdout
 			stdoutConnectorPair = &vm.IOConnectorPair{
-				ReadConnector:  vm.FIFOConnector(fifoSet.Stdout),
+				ReadConnector:  vm.ReadFIFOConnector(fifoSet.Stdout),
 				WriteConnector: vm.VSockAcceptConnector(extraData.StdoutPort),
 			}
 			ts.addCleanup(taskExecID, func() error {
@@ -477,7 +477,7 @@ func (ts *TaskService) Exec(requestCtx context.Context, req *taskAPI.ExecProcess
 		if req.Stderr != "" {
 			req.Stderr = fifoSet.Stderr
 			stderrConnectorPair = &vm.IOConnectorPair{
-				ReadConnector:  vm.FIFOConnector(fifoSet.Stderr),
+				ReadConnector:  vm.ReadFIFOConnector(fifoSet.Stderr),
 				WriteConnector: vm.VSockAcceptConnector(extraData.StderrPort),
 			}
 			ts.addCleanup(taskExecID, func() error {

internal/vm/fifo.go

@@ -21,16 +21,16 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-// FIFOConnector adapts containerd's fifo package to the IOConnector interface
-func FIFOConnector(path string) IOConnector {
+// fifoConnector adapts containerd's fifo package to the IOConnector interface
+func fifoConnector(path string, flag int) IOConnector {
 	return func(procCtx context.Context, logger *logrus.Entry) <-chan IOConnectorResult {
 		returnCh := make(chan IOConnectorResult, 1)
 		defer close(returnCh)
 
 		// We open the FIFO synchronously to ensure that the FIFO is created (via O_CREAT) before
-		// it is passed to any task service. O_RDWR ensures that we don't block on the syscall
+		// it is passed to any task service. O_NONBLOCK ensures that we don't block on the syscall
 		// level (as documented in the fifo pkg).
-		fifo, err := fifo.OpenFifo(procCtx, path, syscall.O_CREAT|syscall.O_RDWR, 0300)
+		fifo, err := fifo.OpenFifo(procCtx, path, syscall.O_CREAT|syscall.O_NONBLOCK|flag, 0300)
 		returnCh <- IOConnectorResult{
 			ReadWriteCloser: fifo,
 			Err:             err,
@@ -39,3 +39,13 @@ func FIFOConnector(path string) IOConnector {
 		return returnCh
 	}
 }
+
+// ReadFIFOConnector returns a FIFO which is open for reading
+func ReadFIFOConnector(path string) IOConnector {
+	return fifoConnector(path, syscall.O_RDONLY)
+}
+
+// WriteFIFOConnector returns a FIFO which is open for writing
+func WriteFIFOConnector(path string) IOConnector {
+	return fifoConnector(path, syscall.O_WRONLY)
+}

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)
 			})
@@ -128,7 +131,8 @@ func (connectorPair *IOConnectorPair) proxy(
 		logger.Debug("begin copying io")
 		defer logger.Debug("end copying io")
 
-		_, err := io.CopyBuffer(writer, reader, make([]byte, internal.DefaultBufferSize))
+		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,6 +142,7 @@ func (connectorPair *IOConnectorPair) proxy(
 			}
 			copyDone <- err
 		}
+		defer logClose(logger, reader, writer)
 	}()
 
 	return initDone, copyDone
@@ -172,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))
+}

runtime/service.go

@@ -793,6 +793,46 @@ func (s *service) buildRootDrive(req *proto.CreateVMRequest) []models.Drive {
 	return builder.Build()
 }
 
+func (s *service) newIOProxy(logger *logrus.Entry, stdin, stdout, stderr string, extraData *proto.ExtraData) (vm.IOProxy, error) {
+	var ioConnectorSet vm.IOProxy
+
+	relVSockPath, err := s.jailer.JailPath().FirecrackerVSockRelPath()
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to get relative path to firecracker vsock")
+	}
+
+	if vm.IsAgentOnlyIO(stdout, logger) {
+		ioConnectorSet = vm.NewNullIOProxy()
+	} else {
+		var stdinConnectorPair *vm.IOConnectorPair
+		if stdin != "" {
+			stdinConnectorPair = &vm.IOConnectorPair{
+				ReadConnector:  vm.ReadFIFOConnector(stdin),
+				WriteConnector: vm.VSockDialConnector(relVSockPath, extraData.StdinPort),
+			}
+		}
+
+		var stdoutConnectorPair *vm.IOConnectorPair
+		if stdout != "" {
+			stdoutConnectorPair = &vm.IOConnectorPair{
+				ReadConnector:  vm.VSockDialConnector(relVSockPath, extraData.StdoutPort),
+				WriteConnector: vm.WriteFIFOConnector(stdout),
+			}
+		}
+
+		var stderrConnectorPair *vm.IOConnectorPair
+		if stderr != "" {
+			stderrConnectorPair = &vm.IOConnectorPair{
+				ReadConnector:  vm.VSockDialConnector(relVSockPath, extraData.StderrPort),
+				WriteConnector: vm.WriteFIFOConnector(stderr),
+			}
+		}
+
+		ioConnectorSet = vm.NewIOConnectorProxy(stdinConnectorPair, stdoutConnectorPair, stderrConnectorPair)
+	}
+	return ioConnectorSet, nil
+}
+
 func (s *service) Create(requestCtx context.Context, request *taskAPI.CreateTaskRequest) (*taskAPI.CreateTaskResponse, error) {
 	logger := s.logger.WithField("task_id", request.ID)
 	defer logPanicAndDie(logger)
@@ -863,41 +903,9 @@ func (s *service) Create(requestCtx context.Context, request *taskAPI.CreateTask
 		return nil, err
 	}
 
-	relVSockPath, err := s.jailer.JailPath().FirecrackerVSockRelPath()
+	ioConnectorSet, err := s.newIOProxy(logger, request.Stdin, request.Stdout, request.Stderr, extraData)
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to get relative path to firecracker vsock")
-	}
-
-	var ioConnectorSet vm.IOProxy
-
-	if vm.IsAgentOnlyIO(request.Stdout, logger) {
-		ioConnectorSet = vm.NewNullIOProxy()
-	} else {
-		var stdinConnectorPair *vm.IOConnectorPair
-		if request.Stdin != "" {
-			stdinConnectorPair = &vm.IOConnectorPair{
-				ReadConnector:  vm.FIFOConnector(request.Stdin),
-				WriteConnector: vm.VSockDialConnector(relVSockPath, extraData.StdinPort),
-			}
-		}
-
-		var stdoutConnectorPair *vm.IOConnectorPair
-		if request.Stdout != "" {
-			stdoutConnectorPair = &vm.IOConnectorPair{
-				ReadConnector:  vm.VSockDialConnector(relVSockPath, extraData.StdoutPort),
-				WriteConnector: vm.FIFOConnector(request.Stdout),
-			}
-		}
-
-		var stderrConnectorPair *vm.IOConnectorPair
-		if request.Stderr != "" {
-			stderrConnectorPair = &vm.IOConnectorPair{
-				ReadConnector:  vm.VSockDialConnector(relVSockPath, extraData.StderrPort),
-				WriteConnector: vm.FIFOConnector(request.Stderr),
-			}
-		}
-
-		ioConnectorSet = vm.NewIOConnectorProxy(stdinConnectorPair, stdoutConnectorPair, stderrConnectorPair)
+		return nil, err
 	}
 
 	// override the request with the bundle dir that should be used inside the VM
@@ -993,41 +1001,9 @@ func (s *service) Exec(requestCtx context.Context, req *taskAPI.ExecProcessReque
 		return nil, err
 	}
 
-	relVSockPath, err := s.jailer.JailPath().FirecrackerVSockRelPath()
+	ioConnectorSet, err := s.newIOProxy(logger, req.Stdin, req.Stdout, req.Stderr, extraData)
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to get relative path to firecracker vsock")
-	}
-
-	var ioConnectorSet vm.IOProxy
-
-	if vm.IsAgentOnlyIO(req.Stdout, logger) {
-		ioConnectorSet = vm.NewNullIOProxy()
-	} else {
-		var stdinConnectorPair *vm.IOConnectorPair
-		if req.Stdin != "" {
-			stdinConnectorPair = &vm.IOConnectorPair{
-				ReadConnector:  vm.FIFOConnector(req.Stdin),
-				WriteConnector: vm.VSockDialConnector(relVSockPath, extraData.StdinPort),
-			}
-		}
-
-		var stdoutConnectorPair *vm.IOConnectorPair
-		if req.Stdout != "" {
-			stdoutConnectorPair = &vm.IOConnectorPair{
-				ReadConnector:  vm.VSockDialConnector(relVSockPath, extraData.StdoutPort),
-				WriteConnector: vm.FIFOConnector(req.Stdout),
-			}
-		}
-
-		var stderrConnectorPair *vm.IOConnectorPair
-		if req.Stderr != "" {
-			stderrConnectorPair = &vm.IOConnectorPair{
-				ReadConnector:  vm.VSockDialConnector(relVSockPath, extraData.StderrPort),
-				WriteConnector: vm.FIFOConnector(req.Stderr),
-			}
-		}
-
-		ioConnectorSet = vm.NewIOConnectorProxy(stdinConnectorPair, stdoutConnectorPair, stderrConnectorPair)
+		return nil, err
 	}
 
 	resp, err := s.taskManager.ExecProcess(requestCtx, req, s.agentClient, ioConnectorSet)