From 34f41dfcbffae9af308f3a0093d9c68070be235d Mon Sep 17 00:00:00 2001
From: Nicolas Lacasse <nlacasse@google.com>
Date: Thu, 3 Feb 2022 10:27:05 -0800
Subject: [PATCH] Hold a reference while calling p9.pathNode.removeWithName()
 callback.

This prevents a racing clunk() call from destroying the file while the callback
is running, leading to potential data races.

PiperOrigin-RevId: 426180506
---
 pkg/p9/path_tree.go | 11 ++++++++++-
 pkg/p9/server.go    | 14 ++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/pkg/p9/path_tree.go b/pkg/p9/path_tree.go
index 14f4e1d68..9b779418c 100644
--- a/pkg/p9/path_tree.go
+++ b/pkg/p9/path_tree.go
@@ -216,8 +216,17 @@ func (p *pathNode) removeWithName(name string, fn func(ref *fidRef)) *pathNode {
 		for ref := range m {
 			delete(m, ref)
 			delete(p.childRefNames, ref)
-			if fn != nil {
+			if fn == nil {
+				// No callback provided.
+				continue
+			}
+			// Attempt to hold a reference while calling fn() to
+			// prevent concurrent destruction of the child, which
+			// can lead to data races. If the child has already
+			// been destroyed, then we can skip the callback.
+			if ref.TryIncRef() {
 				fn(ref)
+				ref.DecRef()
 			}
 		}
 	}
diff --git a/pkg/p9/server.go b/pkg/p9/server.go
index fc6d96828..fe7b57da4 100644
--- a/pkg/p9/server.go
+++ b/pkg/p9/server.go
@@ -205,6 +205,20 @@ func (f *fidRef) DecRef() {
 	}
 }
 
+// TryIncRef returns true if a new reference is taken on the fid, and false if
+// the fid has been destroyed.
+func (f *fidRef) TryIncRef() bool {
+	for {
+		r := atomic.LoadInt64(&f.refs)
+		if r <= 0 {
+			return false
+		}
+		if atomic.CompareAndSwapInt64(&f.refs, r, r+1) {
+			return true
+		}
+	}
+}
+
 // isDeleted returns true if this fidRef has been deleted.
 //
 // Precondition: this must be called via safelyRead, safelyWrite or