github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/fd_cache.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  // This file incorporates work covered by the following copyright and
    16  // permission notice:
    17  //
    18  // Copyright 2017 Attic Labs, Inc. All rights reserved.
    19  // Licensed under the Apache License, version 2.0:
    20  // http://www.apache.org/licenses/LICENSE-2.0
    21  
    22  package nbs
    23  
    24  import (
    25  	"os"
    26  	"sort"
    27  	"sync"
    28  )
    29  
    30  func newFDCache(targetSize int) *fdCache {
    31  	return &fdCache{targetSize: targetSize, cache: map[string]fdCacheEntry{}}
    32  }
    33  
    34  // fdCache ref-counts open file descriptors, but doesn't keep a hard cap on
    35  // the number of open files. Once the cache's target size is exceeded, opening
    36  // a new file causes the cache to try to get the cache back to the target size
    37  // by closing fds with zero refs. If there aren't enough such fds, fdCache
    38  // gives up and tries again next time a caller refs a file.
    39  type fdCache struct {
    40  	targetSize int
    41  	mu         sync.Mutex
    42  	cache      map[string]fdCacheEntry
    43  }
    44  
    45  type fdCacheEntry struct {
    46  	refCount uint32
    47  	f        *os.File
    48  }
    49  
    50  // RefFile returns an opened *os.File for the file at |path|, or an error
    51  // indicating why the file could not be opened. If the cache already had an
    52  // entry for |path|, RefFile increments its refcount and returns the cached
    53  // pointer. If not, it opens the file and caches the pointer for others to
    54  // use. If RefFile returns an error, it's guaranteed that no refCounts were
    55  // changed, so it's an error to make a subsequent call to UnrefFile().
    56  // This is intended for clients that hold fds for extremely short periods.
    57  func (fc *fdCache) RefFile(path string) (f *os.File, err error) {
    58  	refFile := func() *os.File {
    59  		if ce, present := fc.cache[path]; present {
    60  			ce.refCount++
    61  			fc.cache[path] = ce
    62  			return ce.f
    63  		}
    64  		return nil
    65  	}
    66  
    67  	f = func() *os.File {
    68  		fc.mu.Lock()
    69  		defer fc.mu.Unlock()
    70  		return refFile()
    71  	}()
    72  	if f != nil {
    73  		return f, nil
    74  	}
    75  
    76  	// Very much want this to be outside the lock, but the downside is that multiple callers may get here concurrently. That means we need to deal with the raciness below.
    77  	f, err = os.Open(path)
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  
    82  	fc.mu.Lock()
    83  	defer fc.mu.Unlock()
    84  	if cached := refFile(); cached != nil {
    85  		// Someone beat us to it, so close f and return cached fd
    86  		f.Close()
    87  		return cached, nil
    88  	}
    89  	// I won the race!
    90  	fc.cache[path] = fdCacheEntry{f: f, refCount: 1}
    91  	return f, nil
    92  }
    93  
    94  // UnrefFile reduces the refcount of the entry at |path|. If the cache is over
    95  // |fc.targetSize|, UnrefFile makes a best effort to shrink the cache by dumping
    96  // entries with a zero refcount. If there aren't enough zero refcount entries
    97  // to drop to get the cache back to |fc.targetSize|, the cache will remain
    98  // over |fc.targetSize| until the next call to UnrefFile().
    99  func (fc *fdCache) UnrefFile(path string) error {
   100  	fc.mu.Lock()
   101  	defer fc.mu.Unlock()
   102  	if ce, present := fc.cache[path]; present {
   103  		ce.refCount--
   104  		fc.cache[path] = ce
   105  	}
   106  	if len(fc.cache) > fc.targetSize {
   107  		// Sadly, we can't remove items from a map while iterating, so we'll record the stuff we want to drop and then do it after
   108  		needed := len(fc.cache) - fc.targetSize
   109  		toDrop := make([]string, 0, needed)
   110  		for p, ce := range fc.cache {
   111  			if ce.refCount != 0 {
   112  				continue
   113  			}
   114  			toDrop = append(toDrop, p)
   115  			err := ce.f.Close()
   116  
   117  			if err != nil {
   118  				return err
   119  			}
   120  
   121  			needed--
   122  			if needed == 0 {
   123  				break
   124  			}
   125  		}
   126  		for _, p := range toDrop {
   127  			delete(fc.cache, p)
   128  		}
   129  	}
   130  
   131  	return nil
   132  }
   133  
   134  // ShrinkCache forcefully removes all file handles with a refcount of zero.
   135  func (fc *fdCache) ShrinkCache() error {
   136  	fc.mu.Lock()
   137  	defer fc.mu.Unlock()
   138  	toDrop := make([]string, 0, len(fc.cache))
   139  	for p, ce := range fc.cache {
   140  		if ce.refCount != 0 {
   141  			continue
   142  		}
   143  		toDrop = append(toDrop, p)
   144  		err := ce.f.Close()
   145  
   146  		if err != nil {
   147  			return err
   148  		}
   149  	}
   150  
   151  	for _, p := range toDrop {
   152  		delete(fc.cache, p)
   153  	}
   154  
   155  	return nil
   156  }
   157  
   158  // Drop dumps the entire cache and closes all currently open files.
   159  func (fc *fdCache) Drop() {
   160  	fc.mu.Lock()
   161  	defer fc.mu.Unlock()
   162  	for _, ce := range fc.cache {
   163  		ce.f.Close()
   164  	}
   165  	fc.cache = map[string]fdCacheEntry{}
   166  }
   167  
   168  // reportEntries is meant for testing.
   169  func (fc *fdCache) reportEntries() sort.StringSlice {
   170  	fc.mu.Lock()
   171  	defer fc.mu.Unlock()
   172  	ret := make(sort.StringSlice, 0, len(fc.cache))
   173  	for p := range fc.cache {
   174  		ret = append(ret, p)
   175  	}
   176  	sort.Sort(ret)
   177  	return ret
   178  }