Line data    Source code

       1             : // Copyright 2024 The LevelDB-Go and Pebble Authors. All rights reserved. Use
       2             : // of this source code is governed by a BSD-style license that can be found in
       3             : // the LICENSE file.
       4             : 
       5             : package vfstest
       6             : 
       7             : import (
       8             :         "fmt"
       9             :         "io"
      10             :         "os"
      11             :         "runtime"
      12             :         "sync"
      13             : 
      14             :         "github.com/cockroachdb/pebble/vfs"
      15             : )
      16             : 
      17             : // WithOpenFileTracking wraps a FS, returning an FS that will monitor open
      18             : // files. The second return value is a func that when invoked prints the stacks
      19             : // that opened the currently open files. If no files are open, the func writes
      20             : // nothing.
      21           1 : func WithOpenFileTracking(inner vfs.FS) (vfs.FS, func(io.Writer)) {
      22           1 :         wrappedFS := &openFilesFS{
      23           1 :                 inner: inner,
      24           1 :                 files: make(map[*openFile]struct{}),
      25           1 :         }
      26           1 :         return wrappedFS, wrappedFS.dumpStacks
      27           1 : }
      28             : 
      29             : type openFilesFS struct {
      30             :         inner vfs.FS
      31             :         mu    sync.Mutex
      32             :         files map[*openFile]struct{}
      33             : }
      34             : 
      35             : var _ vfs.FS = (*openFilesFS)(nil)
      36             : 
      37           1 : func (fs *openFilesFS) dumpStacks(w io.Writer) {
      38           1 :         fs.mu.Lock()
      39           1 :         defer fs.mu.Unlock()
      40           1 :         if len(fs.files) == 0 {
      41           1 :                 return
      42           1 :         }
      43           1 :         fmt.Fprintf(w, "%d open files:\n", len(fs.files))
      44           1 :         for f := range fs.files {
      45           1 :                 f.dumpStack(w)
      46           1 :                 fmt.Fprintln(w)
      47           1 :         }
      48             : }
      49             : 
      50           1 : func (fs *openFilesFS) Create(name string, category vfs.DiskWriteCategory) (vfs.File, error) {
      51           1 :         return fs.wrapOpenFile(fs.inner.Create(name, category))
      52           1 : }
      53             : 
      54           0 : func (fs *openFilesFS) Link(oldname, newname string) error {
      55           0 :         return fs.inner.Link(oldname, newname)
      56           0 : }
      57             : 
      58           1 : func (fs *openFilesFS) Open(name string, opts ...vfs.OpenOption) (vfs.File, error) {
      59           1 :         return fs.wrapOpenFile(fs.inner.Open(name, opts...))
      60           1 : }
      61             : 
      62             : func (fs *openFilesFS) OpenReadWrite(
      63             :         name string, category vfs.DiskWriteCategory, opts ...vfs.OpenOption,
      64           0 : ) (vfs.File, error) {
      65           0 :         return fs.wrapOpenFile(fs.inner.OpenReadWrite(name, category, opts...))
      66           0 : }
      67             : 
      68           1 : func (fs *openFilesFS) OpenDir(name string) (vfs.File, error) {
      69           1 :         return fs.wrapOpenFile(fs.inner.OpenDir(name))
      70           1 : }
      71             : 
      72           1 : func (fs *openFilesFS) Remove(name string) error {
      73           1 :         return fs.inner.Remove(name)
      74           1 : }
      75             : 
      76           0 : func (fs *openFilesFS) RemoveAll(name string) error {
      77           0 :         return fs.inner.RemoveAll(name)
      78           0 : }
      79             : 
      80           0 : func (fs *openFilesFS) Rename(oldname, newname string) error {
      81           0 :         return fs.inner.Rename(oldname, newname)
      82           0 : }
      83             : 
      84             : func (fs *openFilesFS) ReuseForWrite(
      85             :         oldname, newname string, category vfs.DiskWriteCategory,
      86           1 : ) (vfs.File, error) {
      87           1 :         return fs.wrapOpenFile(fs.inner.ReuseForWrite(oldname, newname, category))
      88           1 : }
      89             : 
      90           1 : func (fs *openFilesFS) MkdirAll(dir string, perm os.FileMode) error {
      91           1 :         return fs.inner.MkdirAll(dir, perm)
      92           1 : }
      93             : 
      94           0 : func (fs *openFilesFS) Lock(name string) (io.Closer, error) {
      95           0 :         return fs.inner.Lock(name)
      96           0 : }
      97             : 
      98           1 : func (fs *openFilesFS) List(dir string) ([]string, error) {
      99           1 :         return fs.inner.List(dir)
     100           1 : }
     101             : 
     102           0 : func (fs *openFilesFS) Stat(name string) (vfs.FileInfo, error) {
     103           0 :         return fs.inner.Stat(name)
     104           0 : }
     105             : 
     106           0 : func (fs *openFilesFS) PathBase(path string) string {
     107           0 :         return fs.inner.PathBase(path)
     108           0 : }
     109             : 
     110           1 : func (fs *openFilesFS) PathJoin(elem ...string) string {
     111           1 :         return fs.inner.PathJoin(elem...)
     112           1 : }
     113             : 
     114           0 : func (fs *openFilesFS) PathDir(path string) string {
     115           0 :         return fs.inner.PathDir(path)
     116           0 : }
     117             : 
     118           0 : func (fs *openFilesFS) GetDiskUsage(path string) (vfs.DiskUsage, error) {
     119           0 :         return fs.inner.GetDiskUsage(path)
     120           0 : }
     121             : 
     122           0 : func (fs *openFilesFS) Unwrap() vfs.FS { return fs.inner }
     123             : 
     124           1 : func (fs *openFilesFS) wrapOpenFile(f vfs.File, err error) (vfs.File, error) {
     125           1 :         if f == nil || err != nil {
     126           0 :                 return f, err
     127           0 :         }
     128           1 :         of := &openFile{File: f, parent: fs}
     129           1 :         of.n = runtime.Callers(2, of.pcs[:])
     130           1 :         fs.mu.Lock()
     131           1 :         defer fs.mu.Unlock()
     132           1 :         fs.files[of] = struct{}{}
     133           1 :         return of, nil
     134             : }
     135             : 
     136             : type openFile struct {
     137             :         vfs.File
     138             :         parent *openFilesFS
     139             :         pcs    [20]uintptr
     140             :         n      int
     141             : }
     142             : 
     143           1 : func (f *openFile) dumpStack(w io.Writer) {
     144           1 :         frames := runtime.CallersFrames(f.pcs[:f.n])
     145           1 :         for {
     146           1 :                 frame, more := frames.Next()
     147           1 :                 fmt.Fprintf(w, "%s\n %s:%d\n", frame.Function, frame.File, frame.Line)
     148           1 :                 if !more {
     149           1 :                         break
     150             :                 }
     151             :         }
     152             : }
     153             : 
     154           1 : func (f *openFile) Close() error {
     155           1 :         err := f.File.Close()
     156           1 :         f.parent.mu.Lock()
     157           1 :         defer f.parent.mu.Unlock()
     158           1 :         delete(f.parent.files, f)
     159           1 :         return err
     160           1 : }