bug: filesLockMap grows unboundedly, leaking memory (#2256) (#2286)

fs.go

@@ -1682,12 +1682,13 @@ func (h *fsHandler) compressAndOpenFSFile(filePath, fileEncoding string) (*fsFil
 		return nil, fmt.Errorf("cannot determine absolute path for %q: %v", compressedFilePath, err)
 	}
 
-	flock := getFileLock(absPath)
-	flock.Lock()
-	ff, err := h.compressFileNolock(f, fileInfo, filePath, compressedFilePath, fileEncoding)
-	flock.Unlock()
-
-	return ff, err
+	flock := acquireFileLock(absPath)
+	flock.mu.Lock()
+	defer func() {
+		flock.mu.Unlock()
+		releaseFileLock(absPath, flock)
+	}()
+	return h.compressFileNolock(f, fileInfo, filePath, compressedFilePath, fileEncoding)
 }
 
 func (h *fsHandler) compressFileNolock(
@@ -2048,12 +2049,36 @@ func fsModTime(t time.Time) time.Time {
 	return t.In(time.UTC).Truncate(time.Second)
 }
 
-var filesLockMap sync.Map
+var (
+	filesLockMu  sync.Mutex
+	filesLockMap = make(map[string]*fileLock)
+)
 
-func getFileLock(absPath string) *sync.Mutex {
-	v, _ := filesLockMap.LoadOrStore(absPath, &sync.Mutex{})
-	filelock := v.(*sync.Mutex)
-	return filelock
+type fileLock struct {
+	mu sync.Mutex
+	// refs counts goroutines that hold or are waiting on mu.
+	refs int
+}
+
+func acquireFileLock(absPath string) *fileLock {
+	filesLockMu.Lock()
+	flock := filesLockMap[absPath]
+	if flock == nil {
+		flock = &fileLock{}
+		filesLockMap[absPath] = flock
+	}
+	flock.refs++
+	filesLockMu.Unlock()
+	return flock
+}
+
+func releaseFileLock(absPath string, flock *fileLock) {
+	filesLockMu.Lock()
+	flock.refs--
+	if flock.refs == 0 && filesLockMap[absPath] == flock {
+		delete(filesLockMap, absPath)
+	}
+	filesLockMu.Unlock()
 }
 
 var _ fs.FS = (*osFS)(nil)

fs_test.go

@@ -772,6 +772,106 @@ func TestFSCompressSingleThreadSkipCache(t *testing.T) {
 	})
 }
 
+func TestFileLockRefCountUsesPathLock(t *testing.T) {
+	path := "/tmp/" + t.Name()
+	otherPath := path + "-other"
+	t.Cleanup(func() {
+		filesLockMu.Lock()
+		delete(filesLockMap, path)
+		delete(filesLockMap, otherPath)
+		filesLockMu.Unlock()
+	})
+
+	firstLock := acquireFileLock(path)
+	secondLock := acquireFileLock(path)
+	if firstLock != secondLock {
+		t.Fatalf("same path must return same lock")
+	}
+
+	otherLock := acquireFileLock(otherPath)
+	if firstLock == otherLock {
+		t.Fatalf("different paths must not share locks")
+	}
+
+	releaseFileLock(path, firstLock)
+	filesLockMu.Lock()
+	remainingRefs := 0
+	if lock := filesLockMap[path]; lock != nil {
+		remainingRefs = lock.refs
+	}
+	filesLockMu.Unlock()
+	if remainingRefs != 1 {
+		t.Fatalf("unexpected remaining refs %d. Expecting 1", remainingRefs)
+	}
+
+	releaseFileLock(path, secondLock)
+	filesLockMu.Lock()
+	_, ok := filesLockMap[path]
+	filesLockMu.Unlock()
+	if ok {
+		t.Fatalf("lock was not removed")
+	}
+
+	releaseFileLock(otherPath, otherLock)
+	filesLockMu.Lock()
+	_, ok = filesLockMap[otherPath]
+	filesLockMu.Unlock()
+	if ok {
+		t.Fatalf("other lock was not removed")
+	}
+}
+
+func TestFileLockWaiterKeepsEntry(t *testing.T) {
+	path := "/tmp/" + t.Name()
+	t.Cleanup(func() {
+		filesLockMu.Lock()
+		delete(filesLockMap, path)
+		filesLockMu.Unlock()
+	})
+
+	firstLock := acquireFileLock(path)
+	firstLock.mu.Lock()
+	waiterReady := make(chan *fileLock, 1)
+	waiterDone := make(chan struct{})
+	go func() {
+		waitingLock := acquireFileLock(path)
+		waiterReady <- waitingLock
+		waitingLock.mu.Lock()
+		defer close(waiterDone)
+		defer releaseFileLock(path, waitingLock)
+		defer waitingLock.mu.Unlock()
+	}()
+
+	waitingLock := <-waiterReady
+	filesLockMu.Lock()
+	currentLock := filesLockMap[path]
+	currentRefs := 0
+	if currentLock != nil {
+		currentRefs = currentLock.refs
+	}
+	filesLockMu.Unlock()
+	if waitingLock != firstLock {
+		t.Errorf("waiter used different lock")
+	}
+	if currentLock != firstLock {
+		t.Errorf("map entry changed while lock had a waiter")
+	}
+	if currentRefs != 2 {
+		t.Errorf("unexpected refs %d. Expecting 2", currentRefs)
+	}
+
+	firstLock.mu.Unlock()
+	releaseFileLock(path, firstLock)
+	<-waiterDone
+
+	filesLockMu.Lock()
+	_, ok := filesLockMap[path]
+	filesLockMu.Unlock()
+	if ok {
+		t.Fatalf("lock was not removed")
+	}
+}
+
 func runFSCompressSingleThread(t *testing.T, fs *FS) {
 	t.Helper()