Merge pull request #657 from austinvazquez/remote-snapshotter-cleaner

add: asynchronous cleanup support to demux snapshotter

Author: austinvazquez

snapshotter/demux/internal/failing_snapshotter.go

@@ -74,3 +74,8 @@ func (s *FailingSnapshotter) Walk(ctx context.Context, fn snapshots.WalkFunc, fi
 func (s *FailingSnapshotter) Close() error {
 	return errors.New("mock Close error from remote snapshotter")
 }
+
+// Cleanup mocks a failing remote call with a non-nil error.
+func (s *FailingSnapshotter) Cleanup(ctx context.Context) error {
+	return errors.New("mock Cleanup error from remote snapshotter")
+}

snapshotter/demux/internal/successful_snapshotter.go

@@ -73,3 +73,8 @@ func (s *SuccessfulSnapshotter) Walk(ctx context.Context, fn snapshots.WalkFunc,
 func (s *SuccessfulSnapshotter) Close() error {
 	return nil
 }
+
+// Cleanup mocks a successful remote call with a nil error.
+func (s *SuccessfulSnapshotter) Cleanup(ctx context.Context) error {
+	return nil
+}

snapshotter/demux/proxy/snapshotter.go

@@ -73,6 +73,14 @@ func NewRemoteSnapshotter(ctx context.Context, address string,
 	return &RemoteSnapshotter{proxy.NewSnapshotter(snapshotsapi.NewSnapshotsClient(gRPCConn), address), metricsProxy}, nil
 }
 
+// Cleanup implements the Cleaner interface for snapshotters.
+// This enables asynchronous resource cleanup by remote snapshotters.
+//
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
+func (rs *RemoteSnapshotter) Cleanup(ctx context.Context) error {
+	return rs.Snapshotter.(snapshots.Cleaner).Cleanup(ctx)
+}
+
 // MetricsProxyPort returns the metrics proxy port for a remote snapshotter.
 func (rs *RemoteSnapshotter) MetricsProxyPort() int {
 	return rs.metricsProxy.Port

snapshotter/demux/snapshotter.go

@@ -27,8 +27,6 @@ import (
 	mountutil "github.com/firecracker-microvm/firecracker-containerd/snapshotter/internal/mount"
 )
 
-const proxiedFunctionCallErrorString = "Proxied function call failed"
-
 // Snapshotter routes snapshotter requests to their destined
 // remote snapshotter via their snapshotter namespace.
 //
@@ -46,234 +44,179 @@ func NewSnapshotter(cache cache.Cache, fetchSnapshotter cache.SnapshotterProvide
 
 // Stat proxies remote snapshotter stat request.
 //
-// See https://github.com/containerd/containerd/blob/main/snapshots/snapshotter.go
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
 func (s *Snapshotter) Stat(ctx context.Context, key string) (snapshots.Info, error) {
 	contextLogger := log.G(ctx).WithField("function", "Stat")
-	namespace, err := getNamespaceFromContext(ctx, contextLogger)
-	if err != nil {
-		return snapshots.Info{}, err
-	}
-	logger := contextLogger.WithField("namespace", namespace)
 
-	snapshotter, err := s.getSnapshotterFromCache(ctx, namespace, logger)
+	snapshotter, err := s.getSnapshotterFromCache(ctx, contextLogger)
 	if err != nil {
 		return snapshots.Info{}, err
 	}
 
-	info, err := snapshotter.Stat(ctx, key)
-	if err != nil {
-		contextLogger.WithError(err).Error(proxiedFunctionCallErrorString)
-		return snapshots.Info{}, err
-	}
-	return info, nil
+	return snapshotter.Stat(ctx, key)
 }
 
 // Update proxies remote snapshotter update request.
 //
-// See https://github.com/containerd/containerd/blob/main/snapshots/snapshotter.go
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
 func (s *Snapshotter) Update(ctx context.Context, info snapshots.Info, fieldpaths ...string) (snapshots.Info, error) {
 	contextLogger := log.G(ctx).WithField("function", "Update")
-	namespace, err := getNamespaceFromContext(ctx, contextLogger)
-	if err != nil {
-		return snapshots.Info{}, err
-	}
-	logger := contextLogger.WithField("namespace", namespace)
 
-	snapshotter, err := s.getSnapshotterFromCache(ctx, namespace, logger)
+	snapshotter, err := s.getSnapshotterFromCache(ctx, contextLogger)
 	if err != nil {
 		return snapshots.Info{}, err
 	}
 
-	updatedInfo, err := snapshotter.Update(ctx, info, fieldpaths...)
-	if err != nil {
-		contextLogger.WithError(err).Error(proxiedFunctionCallErrorString)
-		return snapshots.Info{}, err
-	}
-	return updatedInfo, nil
+	return snapshotter.Update(ctx, info, fieldpaths...)
 }
 
 // Usage proxies remote snapshotter usage request.
 //
-// See https://github.com/containerd/containerd/blob/main/snapshots/snapshotter.go
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
 func (s *Snapshotter) Usage(ctx context.Context, key string) (snapshots.Usage, error) {
 	contextLogger := log.G(ctx).WithField("function", "Usage")
-	namespace, err := getNamespaceFromContext(ctx, contextLogger)
-	if err != nil {
-		return snapshots.Usage{}, err
-	}
-	logger := contextLogger.WithField("namespace", namespace)
 
-	snapshotter, err := s.getSnapshotterFromCache(ctx, namespace, logger)
+	snapshotter, err := s.getSnapshotterFromCache(ctx, contextLogger)
 	if err != nil {
 		return snapshots.Usage{}, err
 	}
 
-	usage, err := snapshotter.Usage(ctx, key)
-	if err != nil {
-		contextLogger.WithError(err).Error(proxiedFunctionCallErrorString)
-		return snapshots.Usage{}, err
-	}
-	return usage, nil
+	return snapshotter.Usage(ctx, key)
 }
 
 // Mounts proxies remote snapshotter mounts request.
 //
-// See https://github.com/containerd/containerd/blob/main/snapshots/snapshotter.go
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
 func (s *Snapshotter) Mounts(ctx context.Context, key string) ([]mount.Mount, error) {
 	contextLogger := log.G(ctx).WithField("function", "Mounts")
-	namespace, err := getNamespaceFromContext(ctx, contextLogger)
-	if err != nil {
-		return []mount.Mount{}, err
-	}
-	logger := contextLogger.WithField("namespace", namespace)
 
-	snapshotter, err := s.getSnapshotterFromCache(ctx, namespace, logger)
+	snapshotter, err := s.getSnapshotterFromCache(ctx, contextLogger)
 	if err != nil {
 		return []mount.Mount{}, err
 	}
 
 	mounts, err := snapshotter.Mounts(ctx, key)
 	if err != nil {
-		contextLogger.WithError(err).Error(proxiedFunctionCallErrorString)
 		return []mount.Mount{}, err
 	}
 	return mountutil.Map(mounts, vm.AddLocalMountIdentifier), nil
 }
 
 // Prepare proxies remote snapshotter prepare request.
 //
-// See https://github.com/containerd/containerd/blob/main/snapshots/snapshotter.go
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
 func (s *Snapshotter) Prepare(ctx context.Context, key string, parent string, opts ...snapshots.Opt) ([]mount.Mount, error) {
 	contextLogger := log.G(ctx).WithField("function", "Prepare")
-	namespace, err := getNamespaceFromContext(ctx, contextLogger)
-	if err != nil {
-		return []mount.Mount{}, err
-	}
-	logger := contextLogger.WithField("namespace", namespace)
 
-	snapshotter, err := s.getSnapshotterFromCache(ctx, namespace, logger)
+	snapshotter, err := s.getSnapshotterFromCache(ctx, contextLogger)
 	if err != nil {
 		return []mount.Mount{}, err
 	}
 
 	mounts, err := snapshotter.Prepare(ctx, key, parent, opts...)
 	if err != nil {
-		contextLogger.WithError(err).Error(proxiedFunctionCallErrorString)
 		return []mount.Mount{}, err
 	}
 	return mountutil.Map(mounts, vm.AddLocalMountIdentifier), nil
 }
 
 // View proxies remote snapshotter view request.
 //
-// See https://github.com/containerd/containerd/blob/main/snapshots/snapshotter.go
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
 func (s *Snapshotter) View(ctx context.Context, key string, parent string, opts ...snapshots.Opt) ([]mount.Mount, error) {
 	contextLogger := log.G(ctx).WithField("function", "View")
-	namespace, err := getNamespaceFromContext(ctx, contextLogger)
-	if err != nil {
-		return []mount.Mount{}, err
-	}
-	logger := contextLogger.WithField("namespace", namespace)
 
-	snapshotter, err := s.getSnapshotterFromCache(ctx, namespace, logger)
+	snapshotter, err := s.getSnapshotterFromCache(ctx, contextLogger)
 	if err != nil {
 		return []mount.Mount{}, err
 	}
 
 	mounts, err := snapshotter.View(ctx, key, parent, opts...)
 	if err != nil {
-		contextLogger.WithError(err).Error(proxiedFunctionCallErrorString)
 		return []mount.Mount{}, err
 	}
 	return mountutil.Map(mounts, vm.AddLocalMountIdentifier), nil
 }
 
 // Commit proxies remote snapshotter commit request.
 //
-// See https://github.com/containerd/containerd/blob/main/snapshots/snapshotter.go
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
 func (s *Snapshotter) Commit(ctx context.Context, name string, key string, opts ...snapshots.Opt) error {
 	contextLogger := log.G(ctx).WithField("function", "Commit")
-	namespace, err := getNamespaceFromContext(ctx, contextLogger)
-	if err != nil {
-		return err
-	}
-	logger := contextLogger.WithField("namespace", namespace)
 
-	snapshotter, err := s.getSnapshotterFromCache(ctx, namespace, logger)
+	snapshotter, err := s.getSnapshotterFromCache(ctx, contextLogger)
 	if err != nil {
 		return err
 	}
 
-	err = snapshotter.Commit(ctx, name, key, opts...)
-	if err != nil {
-		contextLogger.WithError(err).Error(proxiedFunctionCallErrorString)
-		return err
-	}
-	return nil
+	return snapshotter.Commit(ctx, name, key, opts...)
 }
 
 // Remove proxies remote snapshotter remove request.
 //
-// See https://github.com/containerd/containerd/blob/main/snapshots/snapshotter.go
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
 func (s *Snapshotter) Remove(ctx context.Context, key string) error {
 	contextLogger := log.G(ctx).WithField("function", "Remove")
-	namespace, err := getNamespaceFromContext(ctx, contextLogger)
-	if err != nil {
-		return err
-	}
-	logger := contextLogger.WithField("namespace", namespace)
 
-	snapshotter, err := s.getSnapshotterFromCache(ctx, namespace, logger)
+	snapshotter, err := s.getSnapshotterFromCache(ctx, contextLogger)
 	if err != nil {
 		return err
 	}
 
-	err = snapshotter.Remove(ctx, key)
-	if err != nil {
-		contextLogger.WithError(err).Error(proxiedFunctionCallErrorString)
-		return err
-	}
-	return nil
+	return snapshotter.Remove(ctx, key)
 }
 
 // Walk proxies remote snapshotter walk request.
 //
-// See https://github.com/containerd/containerd/blob/main/snapshots/snapshotter.go
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
 func (s *Snapshotter) Walk(ctx context.Context, fn snapshots.WalkFunc, filters ...string) error {
 	contextLogger := log.G(ctx).WithField("function", "Walk")
-	namespace, err := getNamespaceFromContext(ctx, contextLogger)
+
+	_, err := getNamespaceFromContext(ctx, contextLogger)
 	if err != nil {
 		contextLogger.Debug("no namespace found, proxying walk function to all cached snapshotters")
 		return s.cache.WalkAll(ctx, fn, filters...)
 	}
-	logger := contextLogger.WithField("namespace", namespace)
 
-	snapshotter, err := s.getSnapshotterFromCache(ctx, namespace, logger)
+	snapshotter, err := s.getSnapshotterFromCache(ctx, contextLogger)
 	if err != nil {
 		return err
 	}
 
-	err = snapshotter.Walk(ctx, fn, filters...)
-	if err != nil {
-		contextLogger.WithError(err).Error(proxiedFunctionCallErrorString)
-		return err
-	}
-	return nil
+	return snapshotter.Walk(ctx, fn, filters...)
 }
 
 // Close calls close on all cached remote snapshotters.
 //
-// See https://github.com/containerd/containerd/blob/main/snapshots/snapshotter.go
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
 func (s *Snapshotter) Close() error {
 	return s.cache.Close()
 }
 
+// Cleanup proxies remote snapshotter cleanup request.
+//
+// See https://github.com/containerd/containerd/blob/v1.6.4/snapshots/snapshotter.go
+func (s *Snapshotter) Cleanup(ctx context.Context) error {
+	contextLogger := log.G(ctx).WithField("function", "Cleanup")
+
+	snapshotter, err := s.getSnapshotterFromCache(ctx, contextLogger)
+	if err != nil {
+		return err
+	}
+
+	return snapshotter.(snapshots.Cleaner).Cleanup(ctx)
+}
+
 const snapshotterNotFoundErrorString = "Snapshotter not found in cache"
 
-func (s *Snapshotter) getSnapshotterFromCache(ctx context.Context, namespace string, log *logrus.Entry) (snapshots.Snapshotter, error) {
+func (s *Snapshotter) getSnapshotterFromCache(ctx context.Context, log *logrus.Entry) (snapshots.Snapshotter, error) {
+	namespace, err := getNamespaceFromContext(ctx, log)
+	if err != nil {
+		return nil, err
+	}
 	snapshotter, err := s.cache.Get(ctx, namespace, s.fetchSnapshotter)
 	if err != nil {
-		log.WithError(err).Error(snapshotterNotFoundErrorString)
+		log.WithField("namespace", namespace).WithError(err).Error(snapshotterNotFoundErrorString)
 		return nil, err
 	}
 	return snapshotter, nil

snapshotter/demux/snapshotter_test.go

@@ -73,6 +73,7 @@ func TestReturnErrorWhenCalledWithoutNamespacedContext(t *testing.T) {
 		{"View", func() error { _, err := uut.View(ctx, "layerKey", ""); return err }},
 		{"Commit", func() error { return uut.Commit(ctx, "layer1", "layerKey") }},
 		{"Remove", func() error { return uut.Remove(ctx, "layerKey") }},
+		{"Cleanup", func() error { return uut.(snapshots.Cleaner).Cleanup(ctx) }},
 	}
 
 	for _, test := range tests {
@@ -133,6 +134,7 @@ func TestReturnErrorWhenSnapshotterNotFound(t *testing.T) {
 			var callback = func(c context.Context, i snapshots.Info) error { return nil }
 			return uut.Walk(ctx, callback)
 		}},
+		{"Cleanup", func() error { return uut.(snapshots.Cleaner).Cleanup(ctx) }},
 	}
 
 	for _, test := range tests {
@@ -169,6 +171,7 @@ func TestReturnErrorAfterProxyFunctionFailure(t *testing.T) {
 			return uut.Walk(ctx, callback)
 		}},
 		{"Close", func() error { return uut.Close() }},
+		{"Cleanup", func() error { return uut.(snapshots.Cleaner).Cleanup(ctx) }},
 	}
 
 	for _, test := range tests {
@@ -207,6 +210,7 @@ func TestNoErrorIsReturnedOnSuccessfulProxyExecution(t *testing.T) {
 			return uut.Walk(ctx, callback)
 		}},
 		{"Close", func() error { return uut.Close() }},
+		{"Cleanup", func() error { return uut.(snapshots.Cleaner).Cleanup(ctx) }},
 	}
 
 	for _, test := range tests {