Repository: firecracker-microvm/firecracker-containerd
Pull request #339 (2019-11-21 09:58:50)

Support for container stub drive reuse
Signed-off-by: bpopovschi <zyqsempai@mail.ru>

Issue #243 

Added possibility to reuse stub drives of just deleted containers:
- The `Unmount()` method was added to `agent/drive_handler` to support unmounting used stub drives.
- The `Release()` method was added to aggregate all steps we need to release used stub.
- Modified task `Delete()` method to use new `Release()` method.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.


Zyqsempai: @sipsma PTAL and maybe you will have an idea how to fix test, it looks like `hostPath` is missing, but I have no clue how to set it in that test.
Zyqsempai: @sipsma PTAL

Patch:
diff --git a/agent/drive_handler.go b/agent/drive_handler.go
index aeec3ac..c83d383 100644
--- a/agent/drive_handler.go
+++ b/agent/drive_handler.go
@@ -177,22 +177,11 @@ func (dh driveHandler) MountDrive(ctx context.Context, req *drivemount.MountDriv
 	logger = logger.WithField("drive_path", drive.Path())
 
 	// Do a basic check that we won't be mounting over any important system directories
-	resolvedDest, err := evalAnySymlinks(req.DestinationPath)
-	if err != nil {
-		return nil, errors.Wrapf(err,
-			"failed to evaluate any symlinks in drive mount destination %q", req.DestinationPath)
-	}
-
-	for _, systemDir := range bannedSystemDirs {
-		if isOrUnderDir(resolvedDest, systemDir) {
-			return nil, errors.Errorf(
-				"drive mount destination %q resolves to path %q under banned system directory %q",
-				req.DestinationPath, resolvedDest, systemDir,
-			)
-		}
+	if err := isSystemDir(req.DestinationPath); err != nil {
+		return nil, err
 	}
 
-	err = os.MkdirAll(req.DestinationPath, 0700)
+	err := os.MkdirAll(req.DestinationPath, 0700)
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to create drive mount destination %q", req.DestinationPath)
 	}
@@ -229,6 +218,40 @@ func (dh driveHandler) MountDrive(ctx context.Context, req *drivemount.MountDriv
 		drive.Path(), req.DestinationPath)
 }
 
+func (dh driveHandler) UnmountDrive(ctx context.Context, req *drivemount.UnmountDriveRequest) (*empty.Empty, error) {
+	drive, ok := dh.GetDrive(req.DriveID)
+	if !ok {
+		return nil, fmt.Errorf("drive %q could not be found", req.DriveID)
+	}
+
+	err := mount.Unmount(drive.Path(), 0)
+	if err == nil {
+		return &empty.Empty{}, nil
+	}
+
+	return nil, errors.Errorf("failed to unmount the drive %q",
+		drive.Path())
+}
+
+func isSystemDir(path string) error {
+	resolvedDest, err := evalAnySymlinks(path)
+	if err != nil {
+		return errors.Wrapf(err,
+			"failed to evaluate any symlinks in drive ummount destination %q", path)
+	}
+
+	for _, systemDir := range bannedSystemDirs {
+		if isOrUnderDir(resolvedDest, systemDir) {
+			return errors.Errorf(
+				"drive mount destination %q resolves to path %q under banned system directory %q",
+				path, resolvedDest, systemDir,
+			)
+		}
+	}
+
+	return nil
+}
+
 // evalAnySymlinks is similar to filepath.EvalSymlinks, except it will not return an error if part of the
 // provided path does not exist. It will evaluate symlinks present in the path up to a component that doesn't
 // exist, at which point it will just append the rest of the provided path to what has been resolved so far.
diff --git a/proto/service/drivemount/drivemount.proto b/proto/service/drivemount/drivemount.proto
index 5ebe020..4451b21 100644
--- a/proto/service/drivemount/drivemount.proto
+++ b/proto/service/drivemount/drivemount.proto
@@ -6,11 +6,16 @@ option go_package = "drivemount";
 
 service DriveMounter {
     rpc MountDrive(MountDriveRequest) returns (google.protobuf.Empty);
+    rpc UnmountDrive(UnmountDriveRequest) returns (google.protobuf.Empty);
 }
 
 message MountDriveRequest {
+     string DriveID = 1;
+     string DestinationPath = 2;
+     string FilesytemType = 3;
+     repeated string Options = 4;
+}
+
+message UnmountDriveRequest {
     string DriveID = 1;
-    string DestinationPath = 2;
-    string FilesytemType = 3;
-    repeated string Options = 4;
 }
\ No newline at end of file
diff --git a/proto/service/drivemount/ttrpc/drivemount.pb.go b/proto/service/drivemount/ttrpc/drivemount.pb.go
index 2beb99f..3107dee 100644
--- a/proto/service/drivemount/ttrpc/drivemount.pb.go
+++ b/proto/service/drivemount/ttrpc/drivemount.pb.go
@@ -69,14 +69,54 @@ func (m *MountDriveRequest) XXX_DiscardUnknown() {
 
 var xxx_messageInfo_MountDriveRequest proto.InternalMessageInfo
 
+type UnmountDriveRequest struct {
+	DriveID              string   `protobuf:"bytes,1,opt,name=DriveID,json=driveID,proto3" json:"DriveID,omitempty"`
+	XXX_NoUnkeyedLiteral struct{} `json:"-"`
+	XXX_unrecognized     []byte   `json:"-"`
+	XXX_sizecache        int32    `json:"-"`
+}
+
+func (m *UnmountDriveRequest) Reset()      { *m = UnmountDriveRequest{} }
+func (*UnmountDriveRequest) ProtoMessage() {}
+func (*UnmountDriveRequest) Descriptor() ([]byte, []int) {
+	return fileDescriptor_48455e9fea9a7da5, []int{1}
+}
+func (m *UnmountDriveRequest) XXX_Unmarshal(b []byte) error {
+	return m.Unmarshal(b)
+}
+func (m *UnmountDriveRequest) XXX_Marshal(b []byte, deterministic bool) ([]byte, error) {
+	if deterministic {
+		return xxx_messageInfo_UnmountDriveRequest.Marshal(b, m, deterministic)
+	} else {
+		b = b[:cap(b)]
+		n, err := m.MarshalToSizedBuffer(b)
+		if err != nil {
+			return nil, err
+		}
+		return b[:n], nil
+	}
+}
+func (m *UnmountDriveRequest) XXX_Merge(src proto.Message) {
+	xxx_messageInfo_UnmountDriveRequest.Merge(m, src)
+}
+func (m *UnmountDriveRequest) XXX_Size() int {
+	return m.Size()
+}
+func (m *UnmountDriveRequest) XXX_DiscardUnknown() {
+	xxx_messageInfo_UnmountDriveRequest.DiscardUnknown(m)
+}
+
+var xxx_messageInfo_UnmountDriveRequest proto.InternalMessageInfo
+
 func init() {
 	proto.RegisterType((*MountDriveRequest)(nil), "MountDriveRequest")
+	proto.RegisterType((*UnmountDriveRequest)(nil), "UnmountDriveRequest")
 }
 
 func init() { proto.RegisterFile("drivemount.proto", fileDescriptor_48455e9fea9a7da5) }
 
 var fileDescriptor_48455e9fea9a7da5 = []byte{
-	// 240 bytes of a gzipped FileDescriptorProto
+	// 265 bytes of a gzipped FileDescriptorProto
 	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0xe2, 0x12, 0x48, 0x29, 0xca, 0x2c,
 	0x4b, 0xcd, 0xcd, 0x2f, 0xcd, 0x2b, 0xd1, 0x2b, 0x28, 0xca, 0x2f, 0xc9, 0x97, 0x92, 0x4e, 0xcf,
 	0xcf, 0x4f, 0xcf, 0x49, 0xd5, 0x07, 0xf3, 0x92, 0x4a, 0xd3, 0xf4, 0x53, 0x73, 0x0b, 0x4a, 0x2a,
@@ -86,12 +126,14 @@ var fileDescriptor_48455e9fea9a7da5 = []byte{
 	0x2f, 0xb1, 0x24, 0x33, 0x3f, 0x2f, 0x20, 0xb1, 0x24, 0x43, 0x82, 0x09, 0xac, 0x82, 0x3f, 0x05,
 	0x55, 0x58, 0x48, 0x85, 0x8b, 0xd7, 0x2d, 0x33, 0x27, 0xb5, 0xb8, 0xb2, 0x24, 0x35, 0x37, 0xa4,
 	0xb2, 0x20, 0x55, 0x82, 0x19, 0xac, 0x8e, 0x37, 0x0d, 0x59, 0x10, 0x64, 0x93, 0x7f, 0x01, 0x48,
-	0x4f, 0xb1, 0x04, 0x8b, 0x02, 0x33, 0xc8, 0xa6, 0x7c, 0x08, 0xd7, 0xc8, 0x83, 0x8b, 0x07, 0xec,
-	0x06, 0xb0, 0xeb, 0x52, 0x8b, 0x84, 0x2c, 0xb8, 0xb8, 0x10, 0x0e, 0x15, 0x12, 0xd2, 0xc3, 0x70,
-	0xb5, 0x94, 0x98, 0x1e, 0xc4, 0xa7, 0x7a, 0x30, 0x9f, 0xea, 0xb9, 0x82, 0x7c, 0xea, 0xa4, 0x72,
-	0xe2, 0xa1, 0x1c, 0xc3, 0x8d, 0x87, 0x72, 0x0c, 0x0d, 0x8f, 0xe4, 0x18, 0x4f, 0x3c, 0x92, 0x63,
-	0xbc, 0xf0, 0x48, 0x8e, 0xf1, 0xc1, 0x23, 0x39, 0xc6, 0x28, 0x2e, 0x44, 0x60, 0x25, 0xb1, 0x81,
-	0x75, 0x19, 0x03, 0x02, 0x00, 0x00, 0xff, 0xff, 0xc1, 0xad, 0x30, 0xb3, 0x41, 0x01, 0x00, 0x00,
+	0x4f, 0xb1, 0x04, 0x8b, 0x02, 0x33, 0xc8, 0xa6, 0x7c, 0x08, 0x57, 0x49, 0x9f, 0x4b, 0x38, 0x34,
+	0x2f, 0x97, 0x78, 0xa7, 0x19, 0xb5, 0x31, 0x72, 0xf1, 0x80, 0xa5, 0xc0, 0xfe, 0x49, 0x2d, 0x12,
+	0xb2, 0xe0, 0xe2, 0x42, 0x78, 0x4d, 0x48, 0x48, 0x0f, 0xc3, 0x9f, 0x52, 0x62, 0x7a, 0x90, 0xb0,
+	0xd1, 0x83, 0x85, 0x8d, 0x9e, 0x2b, 0x28, 0x6c, 0x84, 0x6c, 0xb8, 0x78, 0x90, 0xed, 0x16, 0x12,
+	0xd1, 0xc3, 0xe2, 0x14, 0x5c, 0xba, 0x9d, 0x54, 0x4e, 0x3c, 0x94, 0x63, 0xb8, 0xf1, 0x50, 0x8e,
+	0xa1, 0xe1, 0x91, 0x1c, 0xe3, 0x89, 0x47, 0x72, 0x8c, 0x17, 0x1e, 0xc9, 0x31, 0x3e, 0x78, 0x24,
+	0xc7, 0x18, 0xc5, 0x85, 0x88, 0x9c, 0x24, 0x36, 0xb0, 0x2e, 0x63, 0x40, 0x00, 0x00, 0x00, 0xff,
+	0xff, 0x1e, 0xd7, 0x06, 0xe8, 0xb1, 0x01, 0x00, 0x00,
 }
 
 func (m *MountDriveRequest) Marshal() (dAtA []byte, err error) {
@@ -151,6 +193,40 @@ func (m *MountDriveRequest) MarshalToSizedBuffer(dAtA []byte) (int, error) {
 	return len(dAtA) - i, nil
 }
 
+func (m *UnmountDriveRequest) Marshal() (dAtA []byte, err error) {
+	size := m.Size()
+	dAtA = make([]byte, size)
+	n, err := m.MarshalToSizedBuffer(dAtA[:size])
+	if err != nil {
+		return nil, err
+	}
+	return dAtA[:n], nil
+}
+
+func (m *UnmountDriveRequest) MarshalTo(dAtA []byte) (int, error) {
+	size := m.Size()
+	return m.MarshalToSizedBuffer(dAtA[:size])
+}
+
+func (m *UnmountDriveRequest) MarshalToSizedBuffer(dAtA []byte) (int, error) {
+	i := len(dAtA)
+	_ = i
+	var l int
+	_ = l
+	if m.XXX_unrecognized != nil {
+		i -= len(m.XXX_unrecognized)
+		copy(dAtA[i:], m.XXX_unrecognized)
+	}
+	if len(m.DriveID) > 0 {
+		i -= len(m.DriveID)
+		copy(dAtA[i:], m.DriveID)
+		i = encodeVarintDrivemount(dAtA, i, uint64(len(m.DriveID)))
+		i--
+		dAtA[i] = 0xa
+	}
+	return len(dAtA) - i, nil
+}
+
 func encodeVarintDrivemount(dAtA []byte, offset int, v uint64) int {
 	offset -= sovDrivemount(v)
 	base := offset
@@ -192,6 +268,22 @@ func (m *MountDriveRequest) Size() (n int) {
 	return n
 }
 
+func (m *UnmountDriveRequest) Size() (n int) {
+	if m == nil {
+		return 0
+	}
+	var l int
+	_ = l
+	l = len(m.DriveID)
+	if l > 0 {
+		n += 1 + l + sovDrivemount(uint64(l))
+	}
+	if m.XXX_unrecognized != nil {
+		n += len(m.XXX_unrecognized)
+	}
+	return n
+}
+
 func sovDrivemount(x uint64) (n int) {
 	return (math_bits.Len64(x|1) + 6) / 7
 }
@@ -212,6 +304,17 @@ func (this *MountDriveRequest) String() string {
 	}, "")
 	return s
 }
+func (this *UnmountDriveRequest) String() string {
+	if this == nil {
+		return "nil"
+	}
+	s := strings.Join([]string{`&UnmountDriveRequest{`,
+		`DriveID:` + fmt.Sprintf("%v", this.DriveID) + `,`,
+		`XXX_unrecognized:` + fmt.Sprintf("%v", this.XXX_unrecognized) + `,`,
+		`}`,
+	}, "")
+	return s
+}
 func valueToStringDrivemount(v interface{}) string {
 	rv := reflect.ValueOf(v)
 	if rv.IsNil() {
@@ -223,6 +326,7 @@ func valueToStringDrivemount(v interface{}) string {
 
 type DriveMounterService interface {
 	MountDrive(ctx context.Context, req *MountDriveRequest) (*empty.Empty, error)
+	UnmountDrive(ctx context.Context, req *UnmountDriveRequest) (*empty.Empty, error)
 }
 
 func RegisterDriveMounterService(srv *github_com_containerd_ttrpc.Server, svc DriveMounterService) {
@@ -234,6 +338,13 @@ func RegisterDriveMounterService(srv *github_com_containerd_ttrpc.Server, svc Dr
 			}
 			return svc.MountDrive(ctx, &req)
 		},
+		"UnmountDrive": func(ctx context.Context, unmarshal func(interface{}) error) (interface{}, error) {
+			var req UnmountDriveRequest
+			if err := unmarshal(&req); err != nil {
+				return nil, err
+			}
+			return svc.UnmountDrive(ctx, &req)
+		},
 	})
 }
 
@@ -254,6 +365,14 @@ func (c *driveMounterClient) MountDrive(ctx context.Context, req *MountDriveRequ
 	}
 	return &resp, nil
 }
+
+func (c *driveMounterClient) UnmountDrive(ctx context.Context, req *UnmountDriveRequest) (*empty.Empty, error) {
+	var resp empty.Empty
+	if err := c.client.Call(ctx, "DriveMounter", "UnmountDrive", req, &resp); err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
 func (m *MountDriveRequest) Unmarshal(dAtA []byte) error {
 	l := len(dAtA)
 	iNdEx := 0
@@ -436,6 +555,92 @@ func (m *MountDriveRequest) Unmarshal(dAtA []byte) error {
 	}
 	return nil
 }
+func (m *UnmountDriveRequest) Unmarshal(dAtA []byte) error {
+	l := len(dAtA)
+	iNdEx := 0
+	for iNdEx < l {
+		preIndex := iNdEx
+		var wire uint64
+		for shift := uint(0); ; shift += 7 {
+			if shift >= 64 {
+				return ErrIntOverflowDrivemount
+			}
+			if iNdEx >= l {
+				return io.ErrUnexpectedEOF
+			}
+			b := dAtA[iNdEx]
+			iNdEx++
+			wire |= uint64(b&0x7F) << shift
+			if b < 0x80 {
+				break
+			}
+		}
+		fieldNum := int32(wire >> 3)
+		wireType := int(wire & 0x7)
+		if wireType == 4 {
+			return fmt.Errorf("proto: UnmountDriveRequest: wiretype end group for non-group")
+		}
+		if fieldNum <= 0 {
+			return fmt.Errorf("proto: UnmountDriveRequest: illegal tag %d (wire type %d)", fieldNum, wire)
+		}
+		switch fieldNum {
+		case 1:
+			if wireType != 2 {
+				return fmt.Errorf("proto: wrong wireType = %d for field DriveID", wireType)
+			}
+			var stringLen uint64
+			for shift := uint(0); ; shift += 7 {
+				if shift >= 64 {
+					return ErrIntOverflowDrivemount
+				}
+				if iNdEx >= l {
+					return io.ErrUnexpectedEOF
+				}
+				b := dAtA[iNdEx]
+				iNdEx++
+				stringLen |= uint64(b&0x7F) << shift
+				if b < 0x80 {
+					break
+				}
+			}
+			intStringLen := int(stringLen)
+			if intStringLen < 0 {
+				return ErrInvalidLengthDrivemount
+			}
+			postIndex := iNdEx + intStringLen
+			if postIndex < 0 {
+				return ErrInvalidLengthDrivemount
+			}
+			if postIndex > l {
+				return io.ErrUnexpectedEOF
+			}
+			m.DriveID = string(dAtA[iNdEx:postIndex])
+			iNdEx = postIndex
+		default:
+			iNdEx = preIndex
+			skippy, err := skipDrivemount(dAtA[iNdEx:])
+			if err != nil {
+				return err
+			}
+			if skippy < 0 {
+				return ErrInvalidLengthDrivemount
+			}
+			if (iNdEx + skippy) < 0 {
+				return ErrInvalidLengthDrivemount
+			}
+			if (iNdEx + skippy) > l {
+				return io.ErrUnexpectedEOF
+			}
+			m.XXX_unrecognized = append(m.XXX_unrecognized, dAtA[iNdEx:iNdEx+skippy]...)
+			iNdEx += skippy
+		}
+	}
+
+	if iNdEx > l {
+		return io.ErrUnexpectedEOF
+	}
+	return nil
+}
 func skipDrivemount(dAtA []byte) (n int, err error) {
 	l := len(dAtA)
 	iNdEx := 0
diff --git a/runtime/drive_handler.go b/runtime/drive_handler.go
index 6ac09bf..a10f401 100644
--- a/runtime/drive_handler.go
+++ b/runtime/drive_handler.go
@@ -96,32 +96,80 @@ type StubDriveHandler struct {
 // Reserve pops a unused stub drive and returns a MountableStubDrive that can be
 // mounted with the provided options as the patched drive information.
 func (h *StubDriveHandler) Reserve(
+	requestCtx context.Context,
 	id string,
 	hostPath string,
 	vmPath string,
 	filesystemType string,
 	options []string,
-) (MountableStubDrive, error) {
+	driveMounter drivemount.DriveMounterService,
+	machine firecracker.MachineIface,
+) error {
 	h.mu.Lock()
 	defer h.mu.Unlock()
 	if len(h.freeDrives) == 0 {
-		return nil, ErrDrivesExhausted
+		return ErrDrivesExhausted
+	}
+	if _, ok := h.usedDrives[id]; ok {
+		// This case means that drive wasn't released or removed properly
+		return fmt.Errorf("drive with ID %s already in use, a previous attempt to remove it may have failed", id)
 	}
 
 	freeDrive := h.freeDrives[0]
 	options, err := setReadWriteOptions(options, freeDrive.driveMount.IsWritable)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	h.freeDrives = h.freeDrives[1:]
-	h.usedDrives[id] = freeDrive
-	return freeDrive.withMountConfig(
+	stubDrive := freeDrive.withMountConfig(
 		hostPath,
 		vmPath,
 		filesystemType,
 		options,
-	), nil
+	)
+	freeDrive = &stubDrive
+
+	err = stubDrive.PatchAndMount(requestCtx, machine, driveMounter)
+	if err != nil {
+		err = errors.Wrapf(err, "failed to mount drive inside vm")
+		return err
+	}
+
+	h.freeDrives = h.freeDrives[1:]
+	h.usedDrives[id] = freeDrive
+	return nil
+}
+
+// Release unmounts stub drive of just deleted container
+// and pushes just released drive to freeDrives
+func (h *StubDriveHandler) Release(
+	requestCtx context.Context,
+	id string,
+	driveMounter drivemount.DriveMounterService,
+	machine firecracker.MachineIface,
+) error {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	drive, ok := h.usedDrives[id]
+	if !ok {
+		return errors.Errorf("container %s drive wasn't found", id)
+	}
+
+	_, err := driveMounter.UnmountDrive(requestCtx, &drivemount.UnmountDriveRequest{
+		DriveID: drive.driveID,
+	})
+	if err != nil {
+		return errors.Wrap(err, "failed to unmount drive")
+	}
+
+	err = machine.UpdateGuestDrive(requestCtx, drive.driveID, filepath.Base(drive.stubPath))
+	if err != nil {
+		return errors.Wrap(err, "failed to patch drive")
+	}
+
+	delete(h.usedDrives, id)
+	h.freeDrives = append(h.freeDrives, drive)
+	return nil
 }
 
 // CreateDriveMountStubs creates a set of MountableStubDrives from the provided DriveMount configs.
diff --git a/runtime/service.go b/runtime/service.go
index c713b25..2b207ca 100644
--- a/runtime/service.go
+++ b/runtime/service.go
@@ -46,6 +46,7 @@ import (
 	"github.com/gofrs/uuid"
 	ptypes "github.com/gogo/protobuf/types"
 	"github.com/golang/protobuf/ptypes/empty"
+	"github.com/hashicorp/go-multierror"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/sys/unix"
@@ -810,21 +811,14 @@ func (s *service) Create(requestCtx context.Context, request *taskAPI.CreateTask
 	}
 	rootfsMnt := request.Rootfs[0]
 
-	stubDrive, err := s.containerStubHandler.Reserve(request.ID,
-		rootfsMnt.Source, vmBundleDir.RootfsPath(), "ext4", nil)
+	err = s.containerStubHandler.Reserve(requestCtx, request.ID,
+		rootfsMnt.Source, vmBundleDir.RootfsPath(), "ext4", nil, s.driveMountClient, s.machine)
 	if err != nil {
 		err = errors.Wrapf(err, "failed to get stub drive for task %q", request.ID)
 		logger.WithError(err).Error()
 		return nil, err
 	}
 
-	err = stubDrive.PatchAndMount(requestCtx, s.machine, s.driveMountClient)
-	if err != nil {
-		err = errors.Wrapf(err, "failed to mount drive inside vm")
-		logger.WithError(err).Error()
-		return nil, err
-	}
-
 	ociConfigBytes, err := hostBundleDir.OCIConfig().Bytes()
 	if err != nil {
 		return nil, err
@@ -928,23 +922,29 @@ func (s *service) Delete(requestCtx context.Context, req *taskAPI.DeleteRequest)
 	if req.ExecID != "" {
 		return resp, nil
 	}
+	var result *multierror.Error
+	// Trying to release stub drive for further reuse
+	err = s.containerStubHandler.Release(requestCtx, req.ID, s.driveMountClient, s.machine)
+	if err != nil {
+		result = multierror.Append(result, errors.Wrapf(err, "failed to release stub drive for container: %s", req.ID))
+	}
 
 	// Otherwise, delete the container
 	dir, err := s.shimDir.BundleLink(req.ID)
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to find the bundle directory of the container: %s", req.ID)
+		result = multierror.Append(result, errors.Wrapf(err, "failed to find the bundle directory of the container: %s", req.ID))
 	}
 
 	_, err = os.Stat(dir.RootPath())
 	if os.IsNotExist(err) {
-		return nil, errors.Wrapf(err, "failed to find the bundle directory of the container: %s", dir.RootPath())
+		result = multierror.Append(result, errors.Wrapf(err, "failed to find the bundle directory of the container: %s", dir.RootPath()))
 	}
 
 	if err = os.Remove(dir.RootPath()); err != nil {
-		return nil, errors.Wrapf(err, "failed to remove the bundle directory of the container: %s", dir.RootPath())
+		result = multierror.Append(result, errors.Wrapf(err, "failed to remove the bundle directory of the container: %s", dir.RootPath()))
 	}
 
-	return resp, nil
+	return resp, result.ErrorOrNil()
 }
 
 // Exec an additional process inside the container


Test patch:
diff --git a/runtime/drive_handler_test.go b/runtime/drive_handler_test.go
index 6a64217..5009402 100644
--- a/runtime/drive_handler_test.go
+++ b/runtime/drive_handler_test.go
@@ -92,18 +92,6 @@ func TestContainerStubs(t *testing.T) {
 		fsType := "foo4"
 		fsOptions := []string{"blah", "bleg"}
 
-		mountableStubDrive, err := stubDriveHandler.Reserve(
-			id, hostPath, vmPath, fsType, fsOptions)
-		assert.NoError(t, err, "failed to reserve stub drive")
-
-		stub := mountableStubDrive.(stubDrive)
-		assert.Equal(t, hostPath, stub.driveMount.HostPath)
-		assert.Equal(t, vmPath, stub.driveMount.VMPath)
-		assert.Equal(t, fsType, stub.driveMount.FilesystemType)
-		assert.Equal(t, append(fsOptions, "rw"), stub.driveMount.Options)
-		assert.True(t, stub.driveMount.IsWritable)
-		assert.Nil(t, stub.driveMount.RateLimiter)
-
 		mockMachine, err := firecracker.NewMachine(ctx, firecracker.Config{}, firecracker.WithClient(
 			firecracker.NewClient("/path/to/socket", nil, false, firecracker.WithOpsClient(&fctesting.MockClient{
 				PatchGuestDriveByIDFn: func(params *ops.PatchGuestDriveByIDParams) (*ops.PatchGuestDriveByIDNoContent, error) {
@@ -120,8 +108,18 @@ func TestContainerStubs(t *testing.T) {
 			expectedOptions:         append(fsOptions, "rw"),
 		}
 
-		err = mountableStubDrive.PatchAndMount(ctx, mockMachine, mockDriveMounter)
-		assert.NoError(t, err, "failed to patch and mount stub drive")
+		err = stubDriveHandler.Reserve(
+			ctx, id, hostPath, vmPath, fsType, fsOptions, mockDriveMounter, mockMachine)
+		assert.NoError(t, err, "failed to reserve stub drive")
+
+		stub, ok := stubDriveHandler.usedDrives[id]
+		assert.True(t, ok)
+		assert.Equal(t, hostPath, stub.driveMount.HostPath)
+		assert.Equal(t, vmPath, stub.driveMount.VMPath)
+		assert.Equal(t, fsType, stub.driveMount.FilesystemType)
+		assert.Equal(t, append(fsOptions, "rw"), stub.driveMount.Options)
+		assert.True(t, stub.driveMount.IsWritable)
+		assert.Nil(t, stub.driveMount.RateLimiter)
 	}
 }
 
diff --git a/runtime/service_integ_test.go b/runtime/service_integ_test.go
index a83352e..74a4d24 100644
--- a/runtime/service_integ_test.go
+++ b/runtime/service_integ_test.go
@@ -782,7 +782,7 @@ func TestCreateContainerWithSameName_Isolated(t *testing.T) {
 	testCreateContainerWithSameName(t, vmID)
 }
 
-func TestCreateTooManyContainers_Isolated(t *testing.T) {
+func TestStubDriveReserveAndReleaseByContainers_Isolated(t *testing.T) {
 	prepareIntegTest(t)
 
 	assert := assert.New(t)
@@ -825,13 +825,9 @@ func TestCreateTooManyContainers_Isolated(t *testing.T) {
 		require.NoError(t, err, "failed to delete a container")
 	}()
 
-	var stdout bytes.Buffer
-	var stderr bytes.Buffer
-
-	// When we reuse a VM explicitly, we cannot start multiple containers unless we pre-allocate stub drives.
-	_, err = c2.NewTask(ctx, cio.NewCreator(cio.WithStreams(nil, &stdout, &stderr)))
-	assert.Contains(err.Error(), "There are no remaining drives to be used")
-	require.Error(t, err)
+	// With the new behaviour, on previous task deletion, stub drive will be released
+	// and now can be reused by new container and task.
+	assert.Equal("hello", startAndWaitTask(ctx, t, c2))
 }
 
 func TestDriveMount_Isolated(t *testing.T) {