Re-establish a connection when vsock breaks

IOProxy struct connects the container's stdio inside a micro VM to
the FIFO files on its host. containerd clients consume the FIFO files
to access the container's stdio.

However, IOProxy was stopping itself when there was an error from
the streams, such as EPIPE. In that case, reading from the FIFO files
blocks indefinitely since there were no writers.

This change adds ioProxyHandler and uses the handler on State() to
make sure IOProxy is running. If not, ioProxyHandler#Attach() creates
a new IOProxy instance.

Fixes firecracker-microvm#482.

Signed-off-by: Kazuyoshi Kato <[email protected]>

Author: kzys

agent/ioproxy_handler.go

@@ -0,0 +1,70 @@
+// 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 main
+
+import (
+	"context"
+
+	"github.com/containerd/containerd/log"
+	"github.com/containerd/containerd/runtime/v2/task"
+	"github.com/firecracker-microvm/firecracker-containerd/internal/vm"
+	ioproxy "github.com/firecracker-microvm/firecracker-containerd/proto/service/ioproxy/ttrpc"
+	"github.com/golang/protobuf/ptypes/empty"
+)
+
+// ioProxyHandler implements IOProxyService that exposes the state of
+// IOProxy instances.
+type ioProxyHandler struct {
+	runcService task.TaskService
+	taskManager vm.TaskManager
+}
+
+var _ ioproxy.IOProxyService = &ioProxyHandler{}
+
+// State returns whether the given exec's IOProxy is still open or not.
+func (ps *ioProxyHandler) State(_ context.Context, req *ioproxy.StateRequest) (*ioproxy.StateResponse, error) {
+	open, err := ps.taskManager.IsProxyOpen(req.ID, req.ExecID)
+	if err != nil {
+		return nil, err
+	}
+	return &ioproxy.StateResponse{IsOpen: open}, nil
+}
+
+// Attach a new IOProxy to the given exec.
+func (ps *ioProxyHandler) Attach(ctx context.Context, req *ioproxy.AttachRequest) (*empty.Empty, error) {
+	state, err := ps.runcService.State(ctx, &task.StateRequest{ID: req.ID, ExecID: req.ExecID})
+	if err != nil {
+		return nil, err
+	}
+
+	logger := log.G(ctx).WithField("TaskID", req.ID).WithField("ExecID", req.ExecID)
+
+	var proxy vm.IOProxy
+	if vm.IsAgentOnlyIO(state.Stdout, logger) {
+		proxy = vm.NewNullIOProxy()
+	} else {
+		proxy = vm.NewIOConnectorProxy(
+			vm.InputPair(req.StdinPort, state.Stdin),
+			vm.OutputPair(state.Stdout, req.StdoutPort),
+			vm.OutputPair(state.Stderr, req.StderrPort),
+		)
+	}
+
+	err = ps.taskManager.AttachIO(ctx, req.ID, req.ExecID, proxy)
+	if err != nil {
+		return nil, err
+	}
+
+	return &empty.Empty{}, nil
+}

agent/main.go

@@ -37,6 +37,7 @@ import (
 	"github.com/firecracker-microvm/firecracker-containerd/internal/vm"
 
 	drivemount "github.com/firecracker-microvm/firecracker-containerd/proto/service/drivemount/ttrpc"
+	ioproxy "github.com/firecracker-microvm/firecracker-containerd/proto/service/ioproxy/ttrpc"
 )
 
 const (
@@ -111,6 +112,11 @@ func main() {
 	}
 	drivemount.RegisterDriveMounterService(server, dh)
 
+	ioproxy.RegisterIOProxyService(server, &ioProxyHandler{
+		runcService: taskService.runcService,
+		taskManager: taskService.taskManager,
+	})
+
 	// Run ttrpc over vsock
 
 	vsockLogger := log.G(shimCtx).WithField("port", port)

agent/service.go

@@ -91,7 +91,7 @@ func NewTaskService(
 	shimCtx context.Context,
 	shimCancel context.CancelFunc,
 	publisher shim.Publisher,
-) (taskAPI.TaskService, error) {
+) (*TaskService, error) {
 	// We provide an empty string for "id" as the service manages multiple tasks; there is no single
 	// "id" being managed. As noted in the comments of the called code, the "id" arg is only used by
 	// the Cleanup function, so it will never be invoked as part of the task service API, which is all

internal/vm/agent.go

@@ -54,3 +54,12 @@ func (*nullIOProxy) start(proc *vmProc) (ioInitDone <-chan error, ioCopyDone <-c
 
 	return initCh, copyCh
 }
+
+// Close is no-op.
+func (*nullIOProxy) Close() {
+}
+
+// IsOpen always returns true, since this proxy implementation is no-op.
+func (*nullIOProxy) IsOpen() bool {
+	return true
+}

internal/vm/ioproxy.go

@@ -17,6 +17,7 @@ import (
 	"context"
 	"io"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/firecracker-microvm/firecracker-containerd/internal"
@@ -36,6 +37,12 @@ type IOProxy interface {
 	// process. It returns two channels, one to indicate io initialization
 	// is completed and one to indicate io copying is completed.
 	start(proc *vmProc) (ioInitDone <-chan error, ioCopyDone <-chan error)
+
+	// Close the proxy.
+	Close()
+
+	// IsOpen returns true if the proxy hasn't been closed.
+	IsOpen() bool
 }
 
 // IOConnector is function that begins initializing an IO connection (i.e.
@@ -71,6 +78,7 @@ func (connectorPair *IOConnectorPair) proxy(
 	logger *logrus.Entry,
 	timeoutAfterExit time.Duration,
 ) (ioInitDone <-chan error, ioCopyDone <-chan error) {
+	// initDone might not have to be buffered. We only send ioInitErr once.
 	initDone := make(chan error, 2)
 	copyDone := make(chan error)
 
@@ -152,6 +160,10 @@ type ioConnectorSet struct {
 	stdin  *IOConnectorPair
 	stdout *IOConnectorPair
 	stderr *IOConnectorPair
+
+	// closeMu is needed since Close() will be called from different goroutines.
+	closeMu sync.Mutex
+	closed  bool
 }
 
 // NewIOConnectorProxy implements the IOProxy interface for a set of
@@ -163,12 +175,33 @@ func NewIOConnectorProxy(stdin, stdout, stderr *IOConnectorPair) IOProxy {
 		stdin:  stdin,
 		stdout: stdout,
 		stderr: stderr,
+		closed: false,
 	}
 }
 
+func (ioConnectorSet *ioConnectorSet) Close() {
+	ioConnectorSet.closeMu.Lock()
+	defer ioConnectorSet.closeMu.Unlock()
+
+	ioConnectorSet.closed = true
+}
+
+func (ioConnectorSet *ioConnectorSet) IsOpen() bool {
+	ioConnectorSet.closeMu.Lock()
+	defer ioConnectorSet.closeMu.Unlock()
+
+	return !ioConnectorSet.closed
+}
+
+// start starts goroutines to copy stdio and returns two channels.
+// The first channel returns its initialization error. The second channel returns errors from copying.
 func (ioConnectorSet *ioConnectorSet) start(proc *vmProc) (ioInitDone <-chan error, ioCopyDone <-chan error) {
 	var initErrG errgroup.Group
-	var copyErrG errgroup.Group
+
+	// When one of the goroutines returns an error, we will cancel
+	// the rest of goroutines through the ctx below.
+	copyErrG, ctx := errgroup.WithContext(proc.ctx)
+
 	waitErrs := func(initErrCh, copyErrCh <-chan error) {
 		initErrG.Go(func() error { return <-initErrCh })
 		copyErrG.Go(func() error { return <-copyErrCh })
@@ -177,24 +210,25 @@ 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.ctx, proc.logger.WithField("stream", "stdin"), 0))
-
+		waitErrs(ioConnectorSet.stdin.proxy(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.ctx, proc.logger.WithField("stream", "stdout"), defaultIOFlushTimeout))
+		waitErrs(ioConnectorSet.stdout.proxy(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.ctx, proc.logger.WithField("stream", "stderr"), defaultIOFlushTimeout))
+		waitErrs(ioConnectorSet.stderr.proxy(ctx, proc.logger.WithField("stream", "stderr"), defaultIOFlushTimeout))
 	} else {
 		proc.logger.Debug("skipping proxy io for unset stderr")
 	}
 
+	// These channels are not buffered, since we will close them right after having one error.
+	// Callers must read the channels.
 	initDone := make(chan error)
 	go func() {
 		defer close(initDone)
@@ -227,3 +261,29 @@ func logClose(logger *logrus.Entry, streams ...io.Closer) {
 		logger.WithError(closeErr).Error("error closing io stream")
 	}
 }
+
+// InputPair returns an IOConnectorPair from the given vsock port to
+// the FIFO file.
+func InputPair(src uint32, dest string) *IOConnectorPair {
+	if dest == "" {
+		return nil
+	}
+
+	return &IOConnectorPair{
+		ReadConnector:  VSockAcceptConnector(src),
+		WriteConnector: WriteFIFOConnector(dest),
+	}
+}
+
+// OutputPair returns an IOConnectorPair from the given FIFO to
+// the vsock port.
+func OutputPair(src string, dest uint32) *IOConnectorPair {
+	if src == "" {
+		return nil
+	}
+
+	return &IOConnectorPair{
+		ReadConnector:  ReadFIFOConnector(src),
+		WriteConnector: VSockAcceptConnector(dest),
+	}
+}