Add optional TTL-based cleanup for httpcache to reduce disk and memory usage

Signed-off-by: abushwang <abushwang@tencent.com>

Author: abushwang

cache/cache.go

@@ -24,6 +24,7 @@ import (
 	"os"
 	"path/filepath"
 	"sync"
+	"time"
 
 	"github.com/containerd/stargz-snapshotter/util/cacheutil"
 	"github.com/containerd/stargz-snapshotter/util/namedmutex"
@@ -65,6 +66,10 @@ type DirectoryCacheConfig struct {
 
 	// FadvDontNeed forcefully clean fscache pagecache for saving memory.
 	FadvDontNeed bool
+
+	// EntryTTL enables TTL-based cleanup for cached blob files.
+	// When zero or negative, TTL-based cleanup is disabled.
+	EntryTTL time.Duration
 }
 
 // TODO: contents validation.
@@ -178,8 +183,12 @@ func NewDirectoryCache(directory string, config DirectoryCacheConfig) (BlobCache
 		bufPool:      bufPool,
 		direct:       config.Direct,
 		fadvDontNeed: config.FadvDontNeed,
+		entryTTL:     config.EntryTTL,
 	}
 	dc.syncAdd = config.SyncAdd
+
+	dc.startCleanupIfNeeded()
+
 	return dc, nil
 }
 
@@ -197,8 +206,11 @@ type directoryCache struct {
 	direct       bool
 	fadvDontNeed bool
 
-	closed   bool
-	closedMu sync.Mutex
+	closed        bool
+	closedMu      sync.Mutex
+	entryTTL      time.Duration
+	cleanupStopCh chan struct{}
+	cleanupWg     sync.WaitGroup
 }
 
 func (dc *directoryCache) Get(key string, opts ...Option) (Reader, error) {
@@ -384,11 +396,19 @@ func (dc *directoryCache) putBuffer(b *bytes.Buffer) {
 
 func (dc *directoryCache) Close() error {
 	dc.closedMu.Lock()
-	defer dc.closedMu.Unlock()
 	if dc.closed {
+		dc.closedMu.Unlock()
 		return nil
 	}
 	dc.closed = true
+	stopCh := dc.cleanupStopCh
+	dc.cleanupStopCh = nil
+	dc.closedMu.Unlock()
+
+	if stopCh != nil {
+		close(stopCh)
+		dc.cleanupWg.Wait()
+	}
 	return os.RemoveAll(dc.directory)
 }
 

cache/httpcache_ttl.go

@@ -0,0 +1,190 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License 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 cache
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"sync"
+	"time"
+
+	"github.com/containerd/log"
+	"github.com/containerd/stargz-snapshotter/util/cacheutil"
+)
+
+func (dc *directoryCache) startCleanupIfNeeded() {
+	if dc.entryTTL <= 0 {
+		return
+	}
+	dc.closedMu.Lock()
+	if dc.closed || dc.cleanupStopCh != nil {
+		dc.closedMu.Unlock()
+		return
+	}
+	stopCh := make(chan struct{})
+	dc.cleanupStopCh = stopCh
+	dc.cleanupWg.Add(1)
+	dc.closedMu.Unlock()
+
+	interval := dc.entryTTL
+	go func() {
+		defer dc.cleanupWg.Done()
+
+		ticker := time.NewTicker(interval)
+		defer ticker.Stop()
+
+		dc.cleanupOnce()
+		for {
+			select {
+			case <-ticker.C:
+				dc.cleanupOnce()
+			case <-stopCh:
+				return
+			}
+		}
+	}()
+}
+
+func (dc *directoryCache) cleanupOnce() {
+	if dc.entryTTL <= 0 {
+		return
+	}
+	if dc.isClosed() {
+		return
+	}
+
+	cutoff := time.Now().Add(-dc.entryTTL)
+	wipBase := dc.wipDirectory
+
+	_ = filepath.WalkDir(dc.directory, func(path string, d fs.DirEntry, walkErr error) error {
+		if walkErr != nil {
+			return nil
+		}
+		if dc.isClosed() {
+			return fs.SkipAll
+		}
+		if d.IsDir() {
+			if path == wipBase {
+				return fs.SkipDir
+			}
+			return nil
+		}
+
+		info, err := d.Info()
+		if err != nil {
+			return nil
+		}
+		if info.ModTime().After(cutoff) {
+			return nil
+		}
+
+		if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
+			log.L.WithError(err).Debugf("failed to remove expired cache entry %q", path)
+		}
+		return nil
+	})
+
+	_ = filepath.WalkDir(wipBase, func(path string, d fs.DirEntry, walkErr error) error {
+		if walkErr != nil {
+			return nil
+		}
+		if dc.isClosed() {
+			return fs.SkipAll
+		}
+		if d.IsDir() {
+			return nil
+		}
+
+		info, err := d.Info()
+		if err != nil {
+			return nil
+		}
+		if info.ModTime().After(cutoff) {
+			return nil
+		}
+		_ = os.Remove(path)
+		return nil
+	})
+}
+
+func NewMemoryCacheWithTTL(ttl time.Duration) BlobCache {
+	if ttl <= 0 {
+		return NewMemoryCache()
+	}
+	bufPool := &sync.Pool{
+		New: func() interface{} {
+			return new(bytes.Buffer)
+		},
+	}
+	c := cacheutil.NewTTLCache(ttl)
+	c.OnEvicted = func(key string, value interface{}) {
+		b := value.(*bytes.Buffer)
+		b.Reset()
+		bufPool.Put(b)
+	}
+	return &ttlMemoryCache{
+		c:       c,
+		bufPool: bufPool,
+	}
+}
+
+type ttlMemoryCache struct {
+	c       *cacheutil.TTLCache
+	bufPool *sync.Pool
+}
+
+func (mc *ttlMemoryCache) Get(key string, opts ...Option) (Reader, error) {
+	v, done, ok := mc.c.Get(key)
+	if !ok {
+		return nil, fmt.Errorf("missed cache: %q", key)
+	}
+	b := v.(*bytes.Buffer)
+	return &reader{
+		ReaderAt: bytes.NewReader(b.Bytes()),
+		closeFunc: func() error {
+			done(false)
+			return nil
+		},
+	}, nil
+}
+
+func (mc *ttlMemoryCache) Add(key string, opts ...Option) (Writer, error) {
+	b := mc.bufPool.Get().(*bytes.Buffer)
+	b.Reset()
+	return &writer{
+		WriteCloser: nopWriteCloser(io.Writer(b)),
+		commitFunc: func() error {
+			mc.c.Remove(key)
+			_, done, _ := mc.c.Add(key, b)
+			done(false)
+			return nil
+		},
+		abortFunc: func() error {
+			b.Reset()
+			mc.bufPool.Put(b)
+			return nil
+		},
+	}, nil
+}
+
+func (mc *ttlMemoryCache) Close() error {
+	return nil
+}

cache/httpcache_ttl_test.go

@@ -0,0 +1,120 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License 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 cache
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+)
+
+func TestNewMemoryCacheWithTTL_Disabled(t *testing.T) {
+	c := NewMemoryCacheWithTTL(0)
+	if _, ok := c.(*MemoryCache); !ok {
+		t.Fatalf("expected *MemoryCache when ttl is disabled; got %T", c)
+	}
+}
+
+func TestNewMemoryCacheWithTTL_Expires(t *testing.T) {
+	ttl := 30 * time.Millisecond
+	c := NewMemoryCacheWithTTL(ttl)
+
+	w, err := c.Add("k1")
+	if err != nil {
+		t.Fatalf("Add failed: %v", err)
+	}
+	if _, err := w.Write([]byte("abc")); err != nil {
+		t.Fatalf("Write failed: %v", err)
+	}
+	if err := w.Commit(); err != nil {
+		t.Fatalf("Commit failed: %v", err)
+	}
+	_ = w.Close()
+
+	r, err := c.Get("k1")
+	if err != nil {
+		t.Fatalf("Get failed: %v", err)
+	}
+	_ = r.Close()
+
+	deadline := time.Now().Add(2 * time.Second)
+	for {
+		if time.Now().After(deadline) {
+			t.Fatalf("entry did not expire within deadline")
+		}
+		time.Sleep(10 * time.Millisecond)
+		if _, err := c.Get("k1"); err == nil {
+			continue
+		}
+		break
+	}
+}
+
+func TestDirectoryCacheCleanupOnce_RemovesExpiredFiles(t *testing.T) {
+	base := t.TempDir()
+	wip := filepath.Join(base, "wip")
+	if err := os.MkdirAll(wip, 0700); err != nil {
+		t.Fatalf("mkdir wip failed: %v", err)
+	}
+
+	dc := &directoryCache{
+		directory:    base,
+		wipDirectory: wip,
+		entryTTL:     100 * time.Millisecond,
+	}
+
+	expired := filepath.Join(base, "aa", "expired")
+	if err := os.MkdirAll(filepath.Dir(expired), 0700); err != nil {
+		t.Fatalf("mkdir expired dir failed: %v", err)
+	}
+	if err := os.WriteFile(expired, []byte("x"), 0600); err != nil {
+		t.Fatalf("write expired failed: %v", err)
+	}
+	old := time.Now().Add(-2 * time.Second)
+	if err := os.Chtimes(expired, old, old); err != nil {
+		t.Fatalf("chtimes expired failed: %v", err)
+	}
+
+	fresh := filepath.Join(base, "bb", "fresh")
+	if err := os.MkdirAll(filepath.Dir(fresh), 0700); err != nil {
+		t.Fatalf("mkdir fresh dir failed: %v", err)
+	}
+	if err := os.WriteFile(fresh, []byte("y"), 0600); err != nil {
+		t.Fatalf("write fresh failed: %v", err)
+	}
+
+	expiredWip := filepath.Join(wip, "tmp-expired")
+	if err := os.WriteFile(expiredWip, []byte("z"), 0600); err != nil {
+		t.Fatalf("write expired wip failed: %v", err)
+	}
+	if err := os.Chtimes(expiredWip, old, old); err != nil {
+		t.Fatalf("chtimes expired wip failed: %v", err)
+	}
+
+	dc.cleanupOnce()
+
+	if _, err := os.Stat(expired); err == nil {
+		t.Fatalf("expected expired file to be removed")
+	}
+	if _, err := os.Stat(expiredWip); err == nil {
+		t.Fatalf("expected expired wip file to be removed")
+	}
+	if _, err := os.Stat(fresh); err != nil {
+		t.Fatalf("expected fresh file to remain; err=%v", err)
+	}
+}

fs/config/config.go

@@ -39,6 +39,11 @@ type Config struct {
 	// Other values default to cache them on disk.
 	HTTPCacheType string `toml:"http_cache_type" json:"http_cache_type"`
 
+	// HTTPCacheChunkTTLSec specifies TTL (in sec) for each http cache chunk.
+	// Default is 0, which disables TTL-based cleanup.
+	// When zero or negative, TTL-based cleanup is disabled.
+	HTTPCacheChunkTTLSec int `toml:"http_cache_chunk_ttl_sec" json:"http_cache_chunk_ttl_sec"`
+
 	// Type of cache for uncompressed files contents. "memory" stores them on memory. Other values
 	// default to cache them on disk.
 	FSCacheType string `toml:"filesystem_cache_type" json:"filesystem_cache_type"`