fix fscache not cleanup

Signed-off-by: abushwang <[email protected]>

Author: abushwang

cache/cache.go

@@ -214,7 +214,7 @@ func (dc *directoryCache) Get(key string, opts ...Option) (Reader, error) {
 			return &reader{
 				ReaderAt: bytes.NewReader(b.(*bytes.Buffer).Bytes()),
 				closeFunc: func() error {
-					done()
+					done(false)
 					return nil
 				},
 			}, nil
@@ -225,7 +225,7 @@ func (dc *directoryCache) Get(key string, opts ...Option) (Reader, error) {
 			return &reader{
 				ReaderAt: f.(*os.File),
 				closeFunc: func() error {
-					done() // file will be closed when it's evicted from the cache
+					done(false) // file will be closed when it's evicted from the cache
 					return nil
 				},
 			}, nil
@@ -257,7 +257,7 @@ func (dc *directoryCache) Get(key string, opts ...Option) (Reader, error) {
 		ReaderAt: file,
 		closeFunc: func() error {
 			_, done, added := dc.fileCache.Add(key, file)
-			defer done() // Release it immediately. Cleaned up on eviction.
+			defer done(false) // Release it immediately. Cleaned up on eviction.
 			if !added {
 				return file.Close() // file already exists in the cache. close it.
 			}
@@ -323,7 +323,7 @@ func (dc *directoryCache) Add(key string, opts ...Option) (Writer, error) {
 				dc.putBuffer(b) // already exists in the cache. abort it.
 			}
 			commit := func() error {
-				defer done()
+				defer done(false)
 				defer w.Close()
 				n, err := w.Write(cached.(*bytes.Buffer).Bytes())
 				if err != nil || n != cached.(*bytes.Buffer).Len() {

fs/fs.go

@@ -1229,7 +1229,9 @@ func (fs *filesystem) Unmount(ctx context.Context, mountpoint string) error {
 	// If the mountpoint is an id-mapped layer, it is pointing to the
 	// underlying layer, so we cannot call done on it.
 	if !isIDMappedDir(mountpoint) {
-		l.Done()
+		if err := l.Close(); err != nil {
+			log.G(ctx).WithError(err).Warn("failed to release resources of the layer")
+		}
 	}
 	fs.layerMu.Unlock()
 	fs.metricsController.Remove(mountpoint)

fs/fs_test.go

@@ -106,4 +106,5 @@ func (l *breakableLayer) Refresh(ctx context.Context, hosts []docker.RegistryHos
 	}
 	return nil
 }
-func (l *breakableLayer) Done() {}
+func (l *breakableLayer) Done()        {}
+func (l *breakableLayer) Close() error { return nil }

fs/layer/layer.go

@@ -104,6 +104,10 @@ 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()
+
+	// Close is the same as Done, but evicts resources related to this layer immediately.
+	// This is used for cleaning up resources on unmount.
+	Close() error
 }
 
 // Info is the current status of a layer.
@@ -247,7 +251,7 @@ func (r *Resolver) Resolve(ctx context.Context, hosts []docker.RegistryHost, ref
 			return &layerRef{l, done}, nil
 		}
 		// Cached layer is invalid
-		done()
+		done(true)
 		r.layerCacheMu.Lock()
 		r.layerCache.Remove(name)
 		r.layerCacheMu.Unlock()
@@ -262,7 +266,7 @@ func (r *Resolver) Resolve(ctx context.Context, hosts []docker.RegistryHost, ref
 	}
 	defer func() {
 		if retErr != nil {
-			blobR.done()
+			blobR.done(true)
 		}
 	}()
 
@@ -364,7 +368,7 @@ func (r *Resolver) resolveBlob(ctx context.Context, hosts []docker.RegistryHost,
 			return &blobRef{blob, done}, nil
 		}
 		// invalid blob. discard this.
-		done()
+		done(true)
 		r.blobCacheMu.Lock()
 		r.blobCache.Remove(name)
 		r.blobCacheMu.Unlock()
@@ -454,7 +458,12 @@ func (l *layer) Refresh(ctx context.Context, hosts []docker.RegistryHost, refspe
 }
 
 func (l *layerRef) Done() {
-	l.done()
+	l.done(false)
+}
+
+func (l *layerRef) Close() error {
+	l.done(true)
+	return nil
 }
 
 func (l *layer) RootNode(baseInode uint32, idMapper idtools.IDMap) (fusefs.InodeEmbedder, error) {
@@ -482,7 +491,7 @@ func (l *layer) close() error {
 	if l.bgResolver != nil {
 		l.bgResolver.Close()
 	}
-	defer l.blob.done() // Close reader first, then close the blob
+	defer l.blob.done(true) // Close reader first, then close the blob
 	return l.r.Close()
 }
 
@@ -511,15 +520,15 @@ func getDisableXAttrAnnotation(desc ocispec.Descriptor) bool {
 // to this blob will be discarded.
 type blobRef struct {
 	remote.Blob
-	done func()
+	done func(bool)
 }
 
 // layerRef is a reference to the layer in the cache. Calling `Done` or `done` decreases the
 // reference counter of this blob in the underlying cache. When nobody refers to the layer in the
 // cache, resources bound to this layer will be discarded.
 type layerRef struct {
 	*layer
-	done func()
+	done func(bool)
 }
 
 type readerAtFunc func([]byte, int64) (int, error)

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/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 {
@@ -430,6 +432,12 @@ 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() {
+		if m.zinfo != nil {
+			m.zinfo.Close()
+		}
+		if m.cache != nil {
+			m.cache.Close()
+		}
+	})
 }

snapshot/snapshot.go

@@ -1117,6 +1117,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

util/lrucache/lrucache.go

@@ -67,7 +67,7 @@ func New(maxEntries int) *Cache {
 // Get retrieves the specified object from the cache and increments the reference counter of the
 // target content. Client must call `done` callback to decrease the reference count when the value
 // will no longer be used.
-func (c *Cache) Get(key string) (value interface{}, done func(), ok bool) {
+func (c *Cache) Get(key string) (value interface{}, done func(bool), ok bool) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	o, ok := c.cache.Get(key)
@@ -83,7 +83,7 @@ func (c *Cache) Get(key string) (value interface{}, done func(), ok bool) {
 // If the specified content already exists in the cache, this sets `added` to false and returns
 // "already cached" content (i.e. doesn't replace the content with the new one). Client must call
 // `done` callback to decrease the counter when the value will no longer be used.
-func (c *Cache) Add(key string, value interface{}) (cachedValue interface{}, done func(), added bool) {
+func (c *Cache) Add(key string, value interface{}) (cachedValue interface{}, done func(bool), added bool) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	if o, ok := c.cache.Get(key); ok {
@@ -110,12 +110,18 @@ func (c *Cache) Remove(key string) {
 	c.cache.Remove(key)
 }
 
-func (c *Cache) decreaseOnceFunc(rc *refCounter) func() {
+func (c *Cache) decreaseOnceFunc(rc *refCounter) func(bool) {
 	var once sync.Once
-	return func() {
+	return func(evict bool) {
 		c.mu.Lock()
 		defer c.mu.Unlock()
 		once.Do(func() { rc.dec() })
+		if evict {
+			rc.finalize()
+			if cached, ok := c.cache.Get(rc.key); ok && cached == rc {
+				c.cache.Remove(rc.key)
+			}
+		}
 	}
 }
 

util/lrucache/lrucache_test.go

@@ -97,13 +97,13 @@ func TestRemove(t *testing.T) {
 		return
 	}
 
-	done1()
+	done1(false)
 	if len(evicted) != 0 {
 		t.Errorf("no content must be evicted until all reference are discarded")
 		return
 	}
 
-	done12()
+	done12(false)
 	if len(evicted) != 1 {
 		t.Errorf("content must be evicted")
 		return
@@ -138,7 +138,7 @@ func TestEviction(t *testing.T) {
 		return
 	}
 
-	done1()
+	done1(false)
 	if len(evicted) != 1 {
 		t.Errorf("1 content must be evicted")
 		return
@@ -148,15 +148,15 @@ func TestEviction(t *testing.T) {
 		return
 	}
 
-	done2() // effective
-	done2() // ignored
-	done2() // ignored
+	done2(false) // effective
+	done2(false) // ignored
+	done2(false) // ignored
 	if len(evicted) != 1 {
 		t.Errorf("only 1 content must be evicted")
 		return
 	}
 
-	done22()
+	done22(false)
 	if len(evicted) != 2 {
 		t.Errorf("2 contents must be evicted")
 		return