From 6037ed1da07d3c7068f19134705e22fbf6189cc6 Mon Sep 17 00:00:00 2001 From: leecha <22087646+leecha@users.noreply.github.com> Date: Thu, 16 Jul 2026 14:04:20 +0800 Subject: [PATCH] fix(diskcache): make Close terminal --- diskcache/README.md | 2 + diskcache/diskcache.go | 8 ++ diskcache/flock_test.go | 10 +- diskcache/flock_unix.go | 20 ++-- diskcache/flock_windows.go | 9 +- diskcache/get.go | 12 +- diskcache/open.go | 52 ++++++--- diskcache/open_test.go | 220 ++++++++++++++++++++++++++++++++++++- diskcache/pos.go | 6 +- diskcache/put.go | 15 +++ diskcache/put_test.go | 122 ++++++++++++++++++++ diskcache/rotate.go | 7 ++ diskcache/write_file.go | 4 + 13 files changed, 449 insertions(+), 38 deletions(-) diff --git a/diskcache/README.md b/diskcache/README.md index 9ada6d77..8570a7d3 100644 --- a/diskcache/README.md +++ b/diskcache/README.md @@ -10,6 +10,8 @@ diskcache 是一种类似 wal 的磁盘缓存，它有如下特性： 限制： - 不支持随机读取，只支持按照 FIFO 的顺序来消费数据 +- `Close()` 是终止操作；关闭后的读写和 rotate 操作会返回 `ErrClosed` +- `.lock` 是持久标记文件；`Close()` 只释放文件锁，不删除该文件 ## 实现算法 diff --git a/diskcache/diskcache.go b/diskcache/diskcache.go index dfc56dd5..4fa4a2c0 100644 --- a/diskcache/diskcache.go +++ b/diskcache/diskcache.go @@ -51,6 +51,9 @@ var ( // Diskcache full, no data can be write now. ErrCacheFull = errors.New("cache full") + // ErrClosed indicates an operation was attempted on a closed cache. + ErrClosed = errors.New("diskcache closed") + ErrInvalidStreamSize = errors.New("invalid stream size") // Invalid cache filename. @@ -72,6 +75,11 @@ type DiskCache struct { dataFiles []string + // lifecycleMu excludes Close from in-flight I/O and protects closed and closeErr. + lifecycleMu sync.RWMutex + closed bool + closeErr error + // current writing/reading file. curWriteFile, curReadfile string diff --git a/diskcache/flock_test.go b/diskcache/flock_test.go index 99525e7e..27d01797 100644 --- a/diskcache/flock_test.go +++ b/diskcache/flock_test.go @@ -16,7 +16,7 @@ import ( ) func TestLockUnlock(t *T.T) { - t.Run("unlock-remove", func(t *T.T) { + t.Run("unlock-keeps-lock-file", func(t *T.T) { p := t.TempDir() fl := newFlock(p) @@ -28,10 +28,10 @@ func TestLockUnlock(t *T.T) { assert.NoError(t, err) t.Logf("fi: %+#v", fi) - fl.unlock() + assert.NoError(t, fl.unlock()) _, err = os.Stat(filepath.Join(p, ".lock")) - assert.Error(t, err) + assert.NoError(t, err) }) t.Run("lock", func(t *T.T) { @@ -48,7 +48,9 @@ func TestLockUnlock(t *T.T) { assert.True(t, ok) assert.NoError(t, err) - defer fl.unlock() + defer func() { + assert.NoError(t, fl.unlock()) + }() time.Sleep(time.Second * 5) }() diff --git a/diskcache/flock_unix.go b/diskcache/flock_unix.go index eacd04f0..90720372 100644 --- a/diskcache/flock_unix.go +++ b/diskcache/flock_unix.go @@ -40,18 +40,22 @@ func (wl *walLock) tryLock() (bool, error) { return true, nil } -func (wl *walLock) unlock() { +func (wl *walLock) unlock() error { if wl.f != nil { - if err := syscall.Flock(int(wl.f.Fd()), syscall.LOCK_UN); err != nil { - l.Errorf("Flock: %s", err.Error()) - } + f := wl.f + wl.f = nil + var errs []error - if err := wl.f.Close(); err != nil { - l.Errorf("CLose: %s", err.Error()) + if err := syscall.Flock(int(f.Fd()), syscall.LOCK_UN); err != nil { + errs = append(errs, fmt.Errorf("unlock lock file: %w", err)) } - if err := os.Remove(wl.file); err != nil { // Optional on Unix - l.Errorf("Remove: %s", err.Error()) + if err := f.Close(); err != nil { + errs = append(errs, fmt.Errorf("close lock file: %w", err)) } + + return errors.Join(errs...) } + + return nil } diff --git a/diskcache/flock_windows.go b/diskcache/flock_windows.go index ce7ffd9f..a35832de 100644 --- a/diskcache/flock_windows.go +++ b/diskcache/flock_windows.go @@ -53,9 +53,12 @@ func (wl *walLock) tryLock() (bool, error) { return true, nil } -func (wl *walLock) unlock() { +func (wl *walLock) unlock() error { if wl.f != nil { - wl.f.Close() // Closing the file handle automatically releases the lock in Windows - os.Remove(wl.file) + f := wl.f + wl.f = nil + return f.Close() // Closing the file handle automatically releases the lock in Windows } + + return nil } diff --git a/diskcache/get.go b/diskcache/get.go index fe40bcb9..4a054e47 100644 --- a/diskcache/get.go +++ b/diskcache/get.go @@ -13,7 +13,8 @@ import ( "time" ) -// Fn is the handler to eat cache from diskcache. +// Fn is the handler to consume data from diskcache. +// A callback must not synchronously call methods on the same DiskCache. type Fn func([]byte) error func (c *DiskCache) switchNextFile() error { @@ -51,6 +52,8 @@ func (c *DiskCache) Get(fn Fn) error { return c.doGet(nil, fn, nil) } +// BufFunc supplies a buffer for BufCallbackGet. +// A callback must not synchronously call methods on the same DiskCache. type BufFunc func() []byte // BufCallbackGet fetch new data from disk cache, and read into buffer that returned by bfn. @@ -70,6 +73,13 @@ func (c *DiskCache) doGet(buf []byte, fn Fn, bfn BufFunc) error { err error ) + c.lifecycleMu.RLock() + defer c.lifecycleMu.RUnlock() + + if c.closed { + return WrapGetError(ErrClosed, c.path, "") + } + c.rlock.Lock() defer c.rlock.Unlock() diff --git a/diskcache/open.go b/diskcache/open.go index 1d2ddc4b..e590c233 100644 --- a/diskcache/open.go +++ b/diskcache/open.go @@ -6,6 +6,7 @@ package diskcache import ( + "errors" "fmt" "os" "path/filepath" @@ -186,42 +187,57 @@ func (c *DiskCache) doOpen() error { return nil } -// Close reclame fd resources. -// Close is safe to call concurrently with other operations and will -// block until all other operations finish. +// Close permanently closes the cache and reclaims its file descriptors. +// It waits for in-flight operations, is idempotent, and causes later I/O +// operations to return ErrClosed. func (c *DiskCache) Close() error { - c.rwlock.Lock() - defer c.rwlock.Unlock() + c.lifecycleMu.Lock() + defer c.lifecycleMu.Unlock() defer func() { lastCloseTimeVec.WithLabelValues(c.path).Set(float64(time.Now().Unix())) }() - if c.rfd != nil { - if err := c.rfd.Close(); err != nil { - return WrapCloseError(err, c.path, "read_fd") - } - c.rfd = nil + if c.closed { + return c.closeErr } + c.closed = true - if !c.noLock { - if c.flock != nil { - c.flock.unlock() + c.rwlock.Lock() + defer c.rwlock.Unlock() + + var errs []error + + if c.rfd != nil { + fd := c.rfd + c.rfd = nil + if err := fd.Close(); err != nil { + errs = append(errs, WrapCloseError(err, c.path, "read_fd")) } } if c.wfd != nil { - if err := c.wfd.Close(); err != nil { - return WrapCloseError(err, c.path, "write_fd") - } + fd := c.wfd c.wfd = nil + if err := fd.Close(); err != nil { + errs = append(errs, WrapCloseError(err, c.path, "write_fd")) + } } if c.pos != nil { if err := c.pos.close(); err != nil { - return WrapPosError(err, c.path, c.pos.Seek).WithDetails("failed_to_close_position_file") + errs = append(errs, WrapPosError(err, c.path, c.pos.Seek).WithDetails("failed_to_close_position_file")) } } - return nil + if !c.noLock && c.flock != nil { + flock := c.flock + c.flock = nil + if err := flock.unlock(); err != nil { + errs = append(errs, NewCacheError(OpUnlock, err, "failed_to_release_directory_lock").WithPath(c.path)) + } + } + + c.closeErr = errors.Join(errs...) + return c.closeErr } diff --git a/diskcache/open_test.go b/diskcache/open_test.go index befd5ec3..847acc91 100644 --- a/diskcache/open_test.go +++ b/diskcache/open_test.go @@ -6,13 +6,16 @@ package diskcache import ( + "errors" "os" "path/filepath" + "strings" "sync" T "testing" "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestOpen(t *T.T) { @@ -26,7 +29,7 @@ func TestOpen(t *T.T) { assert.FileExists(t, filepath.Join(p, ".lock")) assert.NoError(t, c.Close()) - assert.NoFileExists(t, filepath.Join(p, ".lock")) + assert.FileExists(t, filepath.Join(p, ".lock")) c2, err := Open(WithPath(p), WithNoPos(true)) assert.NoError(t, err) @@ -193,3 +196,218 @@ func TestClose(t *T.T) { }) }) } + +func TestClosedCacheRejectsPutAfterOwnershipTransfer(t *T.T) { + p := t.TempDir() + c1, err := Open(WithPath(p), WithNoSync(true), WithNoPos(true)) + require.NoError(t, err) + require.NoError(t, c1.Put([]byte("before-close"))) + require.NoError(t, c1.Close()) + + c2, err := Open(WithPath(p), WithNoSync(true), WithNoPos(true)) + require.NoError(t, err) + t.Cleanup(func() { + assert.NoError(t, c2.Close()) + ResetMetrics() + }) + + require.NoError(t, c2.Put([]byte("current-owner"))) + require.ErrorIs(t, c1.Put([]byte("after-close")), ErrClosed) + require.NoError(t, c2.Rotate()) + + var got [][]byte + for { + err := c2.Get(func(data []byte) error { + got = append(got, append([]byte(nil), data...)) + return nil + }) + if errors.Is(err, ErrNoData) { + break + } + require.NoError(t, err) + } + + require.Equal(t, [][]byte{[]byte("before-close"), []byte("current-owner")}, got) +} + +func TestClosedCacheRejectsOperations(t *T.T) { + t.Cleanup(ResetMetrics) + + tests := []struct { + name string + do func(*DiskCache) error + }{ + { + name: "put", + do: func(c *DiskCache) error { + return c.Put([]byte("data")) + }, + }, + { + name: "stream-put", + do: func(c *DiskCache) error { + return c.StreamPut(strings.NewReader("data"), len("data")) + }, + }, + { + name: "stream-put-invalid-size", + do: func(c *DiskCache) error { + return c.StreamPut(strings.NewReader(""), 0) + }, + }, + { + name: "get", + do: func(c *DiskCache) error { + return c.Get(nil) + }, + }, + { + name: "buf-callback-get", + do: func(c *DiskCache) error { + return c.BufCallbackGet(func() []byte { return make([]byte, 16) }, nil) + }, + }, + { + name: "buf-get", + do: func(c *DiskCache) error { + return c.BufGet(make([]byte, 16), nil) + }, + }, + { + name: "rotate", + do: func(c *DiskCache) error { + return c.Rotate() + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *T.T) { + c, err := Open(WithPath(t.TempDir()), WithNoSync(true)) + require.NoError(t, err) + require.NoError(t, c.Close()) + + dataPath := filepath.Join(c.Path(), "data") + beforeData, err := os.Stat(dataPath) + require.NoError(t, err) + beforeFDs, canCountFDs := countOpenFileDescriptors(c.Path()) + + require.ErrorIs(t, tc.do(c), ErrClosed) + + afterData, err := os.Stat(dataPath) + require.NoError(t, err) + require.Equal(t, beforeData.Size(), afterData.Size()) + if canCountFDs { + afterFDs, ok := countOpenFileDescriptors(c.Path()) + require.True(t, ok) + require.Equal(t, beforeFDs, afterFDs) + } + }) + } +} + +func countOpenFileDescriptors(path string) (int, bool) { + entries, err := os.ReadDir("/proc/self/fd") + if err != nil { + return 0, false + } + + prefix := filepath.Clean(path) + string(os.PathSeparator) + count := 0 + for _, entry := range entries { + target, err := os.Readlink(filepath.Join("/proc/self/fd", entry.Name())) + if err != nil { + continue + } + + target = strings.TrimSuffix(target, " (deleted)") + if strings.HasPrefix(filepath.Clean(target), prefix) { + count++ + } + } + + return count, true +} + +func TestRepeatedCloseDoesNotReleaseNextOwnerLock(t *T.T) { + p := t.TempDir() + c1, err := Open(WithPath(p)) + require.NoError(t, err) + require.NoError(t, c1.Close()) + + c2, err := Open(WithPath(p)) + require.NoError(t, err) + t.Cleanup(func() { + assert.NoError(t, c2.Close()) + ResetMetrics() + }) + + require.NoError(t, c1.Close()) + c3, err := Open(WithPath(p)) + if c3 != nil { + assert.NoError(t, c3.Close()) + } + require.Error(t, err) +} + +func TestCloseReleasesOwnershipAfterDescriptorError(t *T.T) { + p := t.TempDir() + c, err := Open(WithPath(p), WithNoSync(true)) + require.NoError(t, err) + require.NoError(t, c.Put([]byte("data"))) + require.NoError(t, c.Rotate()) + require.NoError(t, c.Get(nil)) + + readFD := c.rfd + writeFD := c.wfd + posFD := c.pos.fd + require.NotNil(t, readFD) + require.NotNil(t, writeFD) + require.NotNil(t, posFD) + require.NoError(t, readFD.Close()) + + t.Cleanup(func() { + _ = writeFD.Close() + _ = posFD.Close() + if c.flock != nil { + _ = c.flock.unlock() + } + ResetMetrics() + }) + + closeErr := c.Close() + require.Error(t, closeErr) + require.EqualError(t, c.Close(), closeErr.Error()) + _, err = writeFD.WriteString("must-be-closed") + require.Error(t, err) + _, err = posFD.WriteString("must-be-closed") + require.Error(t, err) + + c2, err := Open(WithPath(p)) + require.NoError(t, err) + require.NoError(t, c2.Close()) +} + +func TestCloseAggregatesDescriptorErrors(t *T.T) { + p := t.TempDir() + c, err := Open(WithPath(p), WithNoSync(true)) + require.NoError(t, err) + require.NoError(t, c.Put([]byte("data"))) + require.NoError(t, c.Rotate()) + require.NoError(t, c.Get(nil)) + + require.NoError(t, c.rfd.Close()) + require.NoError(t, c.wfd.Close()) + require.NoError(t, c.pos.fd.Close()) + + closeErr := c.Close() + require.Error(t, closeErr) + require.Contains(t, closeErr.Error(), "fd_type=read_fd") + require.Contains(t, closeErr.Error(), "fd_type=write_fd") + require.Contains(t, closeErr.Error(), "failed_to_close_position_file") + + c2, err := Open(WithPath(p)) + require.NoError(t, err) + require.NoError(t, c2.Close()) + ResetMetrics() +} diff --git a/diskcache/pos.go b/diskcache/pos.go index 891fa568..913dde85 100644 --- a/diskcache/pos.go +++ b/diskcache/pos.go @@ -31,12 +31,12 @@ type pos struct { func (p *pos) close() error { if p.fd != nil { - if err := p.fd.Close(); err != nil { + fd := p.fd + p.fd = nil + if err := fd.Close(); err != nil { return WrapFileOperationError(OpClose, err, "", p.fname). WithDetails("failed_to_close_position_fd") } - - p.fd = nil } return nil diff --git a/diskcache/put.go b/diskcache/put.go index bfd47be3..065fbc92 100644 --- a/diskcache/put.go +++ b/diskcache/put.go @@ -24,6 +24,13 @@ func (c *DiskCache) IsFull(newData []byte) bool { func (c *DiskCache) Put(data []byte) error { start := time.Now() // count time before lock + c.lifecycleMu.RLock() + defer c.lifecycleMu.RUnlock() + + if c.closed { + return WrapPutError(ErrClosed, c.path, len(data)) + } + c.wlock.Lock() defer c.wlock.Unlock() @@ -92,6 +99,7 @@ func (c *DiskCache) Put(data []byte) error { // // If we read the data from some network stream(such as HTTP response body), // we can use StreamPut to avoid a intermediate buffer to accept the huge(may be) body. +// The reader must not synchronously call methods on the same DiskCache. func (c *DiskCache) StreamPut(r io.Reader, size int) error { var ( //nolint:ineffassign @@ -101,6 +109,13 @@ func (c *DiskCache) StreamPut(r io.Reader, size int) error { start = time.Now() ) + c.lifecycleMu.RLock() + defer c.lifecycleMu.RUnlock() + + if c.closed { + return NewCacheError(OpStreamPut, ErrClosed, "cache_closed").WithPath(c.path) + } + if size <= 0 { return NewCacheError(OpStreamPut, ErrInvalidStreamSize, fmt.Sprintf("invalid_size=%d", size)).WithPath(c.path) diff --git a/diskcache/put_test.go b/diskcache/put_test.go index 5c3ac126..db1cf77e 100644 --- a/diskcache/put_test.go +++ b/diskcache/put_test.go @@ -7,6 +7,7 @@ package diskcache import ( "errors" + "io" "strings" "sync" "sync/atomic" @@ -482,3 +483,124 @@ func TestStreamPut(t *T.T) { }) }) } + +type gatedReader struct { + reader io.Reader + started chan struct{} + release <-chan struct{} + once sync.Once +} + +func (r *gatedReader) Read(p []byte) (int, error) { + r.once.Do(func() { + close(r.started) + <-r.release + }) + + return r.reader.Read(p) +} + +func TestCloseWaitsForInFlightStreamPut(t *T.T) { + p := t.TempDir() + c, err := Open(WithPath(p), WithNoSync(true), WithNoPos(true)) + require.NoError(t, err) + + started := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + releaseReader := func() { + releaseOnce.Do(func() { + close(release) + }) + } + t.Cleanup(releaseReader) + reader := &gatedReader{ + reader: strings.NewReader("in-flight"), + started: started, + release: release, + } + + putDone := make(chan error, 1) + go func() { + putDone <- c.StreamPut(reader, len("in-flight")) + }() + <-started + + closeDone := make(chan error, 1) + go func() { + closeDone <- c.Close() + }() + requireLifecycleWriter(t, c) + + select { + case err := <-closeDone: + releaseReader() + require.NoError(t, <-putDone) + require.Failf(t, "Close returned before StreamPut finished", "Close error: %v", err) + default: + } + + releaseReader() + require.NoError(t, <-putDone) + require.NoError(t, <-closeDone) + require.ErrorIs(t, c.StreamPut(strings.NewReader("late"), len("late")), ErrClosed) + + ResetMetrics() +} + +func TestCloseWaitsForInFlightPut(t *T.T) { + p := t.TempDir() + c, err := Open(WithPath(p), WithNoSync(true), WithNoPos(true)) + require.NoError(t, err) + t.Cleanup(func() { + _ = c.Close() + ResetMetrics() + }) + + c.wlock.Lock() + wlockHeld := true + t.Cleanup(func() { + if wlockHeld { + c.wlock.Unlock() + } + }) + + putStarted := make(chan struct{}) + putDone := make(chan error, 1) + go func() { + close(putStarted) + putDone <- c.Put([]byte("in-flight")) + }() + <-putStarted + + require.Eventually(t, func() bool { + if c.lifecycleMu.TryLock() { + c.lifecycleMu.Unlock() + return false + } + return true + }, time.Second, time.Millisecond) + + closeDone := make(chan error, 1) + go func() { + closeDone <- c.Close() + }() + requireLifecycleWriter(t, c) + + c.wlock.Unlock() + wlockHeld = false + require.NoError(t, <-putDone) + require.NoError(t, <-closeDone) + require.ErrorIs(t, c.Put([]byte("late")), ErrClosed) +} + +func requireLifecycleWriter(t *T.T, c *DiskCache) { + t.Helper() + require.Eventually(t, func() bool { + if c.lifecycleMu.TryRLock() { + c.lifecycleMu.RUnlock() + return false + } + return true + }, time.Second, time.Millisecond) +} diff --git a/diskcache/rotate.go b/diskcache/rotate.go index 1473525e..05706234 100644 --- a/diskcache/rotate.go +++ b/diskcache/rotate.go @@ -21,6 +21,13 @@ import ( // NOTE: You do not need to call Rotate() during daily usage, we export // that function for testing cases. func (c *DiskCache) Rotate() error { + c.lifecycleMu.RLock() + defer c.lifecycleMu.RUnlock() + + if c.closed { + return NewCacheError(OpRotate, ErrClosed, "cache_closed").WithPath(c.path) + } + return c.rotate() } diff --git a/diskcache/write_file.go b/diskcache/write_file.go index e1c75333..d3a52fbc 100644 --- a/diskcache/write_file.go +++ b/diskcache/write_file.go @@ -24,6 +24,10 @@ func (c *DiskCache) readFileName() string { } func (c *DiskCache) ensureWriteFile() error { + if c.closed { + return ErrClosed + } + if c.wfd != nil { return nil }