fix fscache not cleanup

Signed-off-by: Shubhranshu Mahapatra <shubhum@amazon.com>
Signed-off-by: abushwang <abushwang@tencent.com>

Author: abushwang

fs/fs.go

@@ -1229,8 +1229,13 @@ func (fs *filesystem) Unmount(ctx context.Context, mountpoint string) error {
 	delete(fs.layer, mountpoint)
 	// If the mountpoint is an id-mapped layer, it is pointing to the
 	// underlying layer, so we cannot call done on it.
+	// We do a evict call to call the registered evict functions.
+	// This cleans up layer resources and removes it from the resolver cache,
+	// preparing the layer for removal and freeing associated memory done on
+	// call of the finalizer in span manager
 	if !isIDMappedDir(mountpoint) {
 		l.Done()
+		fs.resolver.Evict(l.GetCacheRefKey())
 	}
 	fs.layerMu.Unlock()
 	fs.metricsController.Remove(mountpoint)

fs/fs_test.go

@@ -95,6 +95,7 @@ func (l *breakableLayer) SkipVerify()                          {}
 func (l *breakableLayer) ReadAt([]byte, int64, ...remote.Option) (int, error) {
 	return 0, fmt.Errorf("fail")
 }
+func (l *breakableLayer) GetCacheRefKey() string { return "" }
 func (l *breakableLayer) BackgroundFetch() error { return fmt.Errorf("fail") }
 func (l *breakableLayer) Check() error {
 	if !l.success {

fs/layer/layer.go

@@ -104,6 +104,9 @@ type Layer interface {
 	// Done releases the reference to this layer. The resources related to this layer will be
 	// discarded sooner or later. Queries after calling this function won't be serviced.
 	Done()
+
+	// GetCacheRefKey returns the reference key for the cache used by the layer
+	GetCacheRefKey() string
 }
 
 // Info is the current status of a layer.
@@ -226,6 +229,12 @@ func newCache(root string, cacheType string, cfg config.FSConfig) (cache.BlobCac
 	)
 }
 
+func (r *Resolver) Evict(name string) {
+	r.layerCacheMu.Lock()
+	r.layerCache.Remove(name)
+	r.layerCacheMu.Unlock()
+}
+
 // Resolve resolves a layer based on the passed layer blob information.
 func (r *Resolver) Resolve(ctx context.Context, hosts []docker.RegistryHost, refspec reference.Spec, desc, sociDesc ocispec.Descriptor, opCounter *FuseOperationCounter, disableVerification bool, metadataOpts ...metadata.Option) (_ Layer, retErr error) {
 	name := refspec.String() + "/" + desc.Digest.String()
@@ -339,7 +348,7 @@ func (r *Resolver) Resolve(ctx context.Context, hosts []docker.RegistryHost, ref
 	}
 	disableXAttrs := getDisableXAttrAnnotation(sociDesc)
 	// Combine layer information together and cache it.
-	l := newLayer(r, desc, blobR, vr, bgLayerResolver, opCounter, disableXAttrs)
+	l := newLayer(r, desc, name, blobR, vr, bgLayerResolver, opCounter, disableXAttrs)
 	r.layerCacheMu.Lock()
 	cachedL, done2, added := r.layerCache.Add(name, l)
 	r.layerCacheMu.Unlock()
@@ -370,18 +379,8 @@ func (r *Resolver) resolveBlob(ctx context.Context, hosts []docker.RegistryHost,
 		r.blobCacheMu.Unlock()
 	}
 
-	httpCache, err := newCache(filepath.Join(r.rootDir, "httpcache"), r.config.HTTPCacheType, r.config)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create http cache: %w", err)
-	}
-	defer func() {
-		if retErr != nil {
-			httpCache.Close()
-		}
-	}()
-
 	// Resolve the blob and cache the result.
-	b, err := r.resolver.Resolve(ctx, hosts, refspec, desc, httpCache)
+	b, err := r.resolver.Resolve(ctx, hosts, refspec, desc)
 	if err != nil {
 		return nil, fmt.Errorf("failed to resolve the source: %w", err)
 	}
@@ -397,6 +396,7 @@ func (r *Resolver) resolveBlob(ctx context.Context, hosts []docker.RegistryHost,
 func newLayer(
 	resolver *Resolver,
 	desc ocispec.Descriptor,
+	cacheRefKey string,
 	blob *blobRef,
 	r reader.Reader,
 	bgResolver backgroundfetcher.Resolver,
@@ -406,6 +406,7 @@ func newLayer(
 	return &layer{
 		resolver:             resolver,
 		desc:                 desc,
+		cacheRefKey:          cacheRefKey,
 		blob:                 blob,
 		r:                    r,
 		bgResolver:           bgResolver,
@@ -415,9 +416,10 @@ func newLayer(
 }
 
 type layer struct {
-	resolver *Resolver
-	desc     ocispec.Descriptor
-	blob     *blobRef
+	resolver    *Resolver
+	desc        ocispec.Descriptor
+	cacheRefKey string
+	blob        *blobRef
 
 	bgResolver backgroundfetcher.Resolver
 
@@ -430,6 +432,10 @@ type layer struct {
 	closedMu sync.Mutex
 }
 
+func (l *layer) GetCacheRefKey() string {
+	return l.cacheRefKey
+}
+
 func (l *layer) Info() Info {
 	return Info{
 		Digest:      l.desc.Digest,

fs/reader/reader.go

@@ -128,6 +128,9 @@ func (gr *reader) Close() (retErr error) {
 		return nil
 	}
 	gr.closed = true
+	if gr.spanManager != nil {
+		gr.spanManager.Close()
+	}
 	if err := gr.r.Close(); err != nil {
 		retErr = errors.Join(retErr, err)
 	}

fs/remote/resolver.go

@@ -97,7 +97,7 @@ func NewResolver(cfg config.BlobConfig, handlers map[string]Handler) *Resolver {
 	}
 }
 
-func (r *Resolver) Resolve(ctx context.Context, hosts []docker.RegistryHost, refspec reference.Spec, desc ocispec.Descriptor, blobCache cache.BlobCache) (Blob, error) {
+func (r *Resolver) Resolve(ctx context.Context, hosts []docker.RegistryHost, refspec reference.Spec, desc ocispec.Descriptor) (Blob, error) {
 
 	var (
 		validInterval = time.Duration(r.blobConfig.ValidInterval) * time.Second

fs/span-manager/span_manager.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"io"
 	"runtime"
+	"sync"
 
 	"github.com/awslabs/soci-snapshotter/cache"
 	"github.com/awslabs/soci-snapshotter/util/ioutils"
@@ -49,6 +50,7 @@ type SpanManager struct {
 	spans                             []*span
 	ztoc                              *ztoc.Ztoc
 	maxSpanVerificationFailureRetries int
+	closeOnce                         sync.Once
 }
 
 type spanInfo struct {
@@ -432,6 +434,13 @@ func (m *SpanManager) verifySpanContents(compressedData []byte, spanID compressi
 
 // Close closes both the underlying zinfo data and blob cache.
 func (m *SpanManager) Close() {
-	m.zinfo.Close()
-	m.cache.Close()
+	m.closeOnce.Do(func() {
+		log.L.Debug("Spancache finalizer called. Cleaning up spancache")
+		if m.zinfo != nil {
+			m.zinfo.Close()
+		}
+		if m.cache != nil {
+			m.cache.Close()
+		}
+	})
 }

snapshot/snapshot.go

@@ -1118,6 +1118,26 @@ func (o *snapshotter) restoreRemoteSnapshot(ctx context.Context) error {
 		return err
 	}
 	for _, info := range task {
+		if err := func() error {
+			ctx, t, err := o.ms.TransactionContext(ctx, false)
+			if err != nil {
+				return err
+			}
+			defer t.Rollback()
+			id, _, _, err := storage.GetInfo(ctx, info.Name)
+			if err != nil {
+				return err
+			}
+			if err := os.Mkdir(filepath.Join(o.root, "snapshots", id), 0700); err != nil && !os.IsExist(err) {
+				return err
+			}
+			if err := os.Mkdir(o.upperPath(id), 0755); err != nil && !os.IsExist(err) {
+				return err
+			}
+			return nil
+		}(); err != nil {
+			return fmt.Errorf("failed to create remote snapshot directory: %s: %w", info.Name, err)
+		}
 		ns, ok := info.Labels[source.TargetNamespace]
 		if !ok {
 			return ErrNoNamespace