Merge pull request firecracker-microvm#483 from kzys/fix-attach-3

Re-establish a connection when vsock breaks

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),
+	}
+}