1 files changed, 159 insertions, 52 deletions
diff --git a/pkg/p9/path_tree.go b/pkg/p9/path_tree.go
index 5a209b291..865459411 100644
--- a/pkg/p9/path_tree.go
+++ b/pkg/p9/path_tree.go
@@ -22,93 +22,200 @@ import (
 // pathNode is a single node in a path traversal.
 //
 // These are shared by all fidRefs that point to the same path.
+//
+// Lock ordering:
+//   opMu
+//     childMu
+//
+//   Two different pathNodes may only be locked if Server.renameMu is held for
+//   write, in which case they can be acquired in any order.
 type pathNode struct {
-	// fidRefName is a map[*fidRef]string mapping child fidRefs to their
-	// path component name.
-	fidRefNames sync.Map
+	// opMu synchronizes high-level, sematic operations, such as the
+	// simultaneous creation and deletion of a file.
+	//
+	// opMu does not directly protect any fields in pathNode.
+	opMu sync.RWMutex
 
-	// childNodes is a map[string]*pathNode mapping child path component
-	// names to their pathNode.
-	childNodes sync.Map
+	// childMu protects the fields below.
+	childMu sync.RWMutex
 
-	// mu does *not* protect the fields above.
-	//
-	// They are not synchronized because we allow certain operations (file
-	// walk) to proceed without having to acquire a write lock. mu exists
-	// to synchronize high-level, semantic operations, such as the
-	// simultaneous creation and deletion of a file.
-	mu sync.RWMutex
+	// childNodes maps child path component names to their pathNode.
+	childNodes map[string]*pathNode
+
+	// childRefs maps child path component names to all of the their
+	// references.
+	childRefs map[string]map[*fidRef]struct{}
+
+	// childRefNames maps child references back to their path component
+	// name.
+	childRefNames map[*fidRef]string
+}
+
+func newPathNode() *pathNode {
+	return &pathNode{
+		childNodes:    make(map[string]*pathNode),
+		childRefs:     make(map[string]map[*fidRef]struct{}),
+		childRefNames: make(map[*fidRef]string),
+	}
+}
+
+// forEachChildRef calls fn for each child reference.
+func (p *pathNode) forEachChildRef(fn func(ref *fidRef, name string)) {
+	p.childMu.RLock()
+	defer p.childMu.RUnlock()
+
+	for name, m := range p.childRefs {
+		for ref := range m {
+			fn(ref, name)
+		}
+	}
+}
+
+// forEachChildNode calls fn for each child pathNode.
+func (p *pathNode) forEachChildNode(fn func(pn *pathNode)) {
+	p.childMu.RLock()
+	defer p.childMu.RUnlock()
+
+	for _, pn := range p.childNodes {
+		fn(pn)
+	}
 }
 
 // pathNodeFor returns the path node for the given name, or a new one.
-//
-// Precondition: This call is synchronized w.r.t. other adding or removing of
-// children.
 func (p *pathNode) pathNodeFor(name string) *pathNode {
-	// Load the existing path node.
-	if pn, ok := p.childNodes.Load(name); ok {
-		return pn.(*pathNode)
+	p.childMu.RLock()
+	// Fast path, node already exists.
+	if pn, ok := p.childNodes[name]; ok {
+		p.childMu.RUnlock()
+		return pn
+	}
+	p.childMu.RUnlock()
+
+	// Slow path, create a new pathNode for shared use.
+	p.childMu.Lock()
+
+	// Re-check after re-lock.
+	if pn, ok := p.childNodes[name]; ok {
+		p.childMu.Unlock()
+		return pn
 	}
 
-	// Create a new pathNode for shared use.
-	pn, _ := p.childNodes.LoadOrStore(name, new(pathNode))
-	return pn.(*pathNode)
+	pn := newPathNode()
+	p.childNodes[name] = pn
+	p.childMu.Unlock()
+	return pn
 }
 
 // nameFor returns the name for the given fidRef.
 //
 // Precondition: addChild is called for ref before nameFor.
 func (p *pathNode) nameFor(ref *fidRef) string {
-	if s, ok := p.fidRefNames.Load(ref); ok {
-		return s.(string)
+	p.childMu.RLock()
+	n, ok := p.childRefNames[ref]
+	p.childMu.RUnlock()
+
+	if !ok {
+		// This should not happen, don't proceed.
+		panic(fmt.Sprintf("expected name for %+v, none found", ref))
 	}
 
-	// This should not happen, don't proceed.
-	panic(fmt.Sprintf("expected name for %+v, none found", ref))
+	return n
 }
 
-// addChild adds a child to p.
+// addChildLocked adds a child reference to p.
 //
-// This applies only to an individual fidRef.
-//
-// Precondition: ref is added only once unless it is removed before adding with
-// a new name.
-func (p *pathNode) addChild(ref *fidRef, name string) {
-	if s, ok := p.fidRefNames.Load(ref); ok {
+// Precondition: As addChild, plus childMu is locked for write.
+func (p *pathNode) addChildLocked(ref *fidRef, name string) {
+	if n, ok := p.childRefNames[ref]; ok {
 		// This should not happen, don't proceed.
-		panic(fmt.Sprintf("unexpected fidRef %+v with path %q, wanted %q", ref, s, name))
+		panic(fmt.Sprintf("unexpected fidRef %+v with path %q, wanted %q", ref, n, name))
 	}
 
-	p.fidRefNames.Store(ref, name)
+	p.childRefNames[ref] = name
+
+	m, ok := p.childRefs[name]
+	if !ok {
+		m = make(map[*fidRef]struct{})
+		p.childRefs[name] = m
+	}
+
+	m[ref] = struct{}{}
+}
+
+// addChild adds a child reference to p.
+//
+// Precondition: ref may only be added once at a time.
+func (p *pathNode) addChild(ref *fidRef, name string) {
+	p.childMu.Lock()
+	p.addChildLocked(ref, name)
+	p.childMu.Unlock()
 }
 
 // removeChild removes the given child.
 //
-// This applies only to an individual fidRef.
+// This applies only to an individual fidRef, which is not required to exist.
 func (p *pathNode) removeChild(ref *fidRef) {
-	p.fidRefNames.Delete(ref)
+	p.childMu.Lock()
+
+	// This ref may not exist anymore. This can occur, e.g., in unlink,
+	// where a removeWithName removes the ref, and then a DecRef on the ref
+	// attempts to remove again.
+	if name, ok := p.childRefNames[ref]; ok {
+		m, ok := p.childRefs[name]
+		if !ok {
+			// This should not happen, don't proceed.
+			p.childMu.Unlock()
+			panic(fmt.Sprintf("name %s missing from childfidRefs", name))
+		}
+
+		delete(m, ref)
+		if len(m) == 0 {
+			delete(p.childRefs, name)
+		}
+	}
+
+	delete(p.childRefNames, ref)
+
+	p.childMu.Unlock()
 }
 
-// removeWithName removes all references with the given name.
+// addPathNodeFor adds an existing pathNode as the node for name.
 //
-// The original pathNode is returned by this function, and removed from this
-// pathNode. Any operations on the removed tree must use this value.
+// Preconditions: newName does not exist.
+func (p *pathNode) addPathNodeFor(name string, pn *pathNode) {
+	p.childMu.Lock()
+
+	if opn, ok := p.childNodes[name]; ok {
+		p.childMu.Unlock()
+		panic(fmt.Sprintf("unexpected pathNode %+v with path %q", opn, name))
+	}
+
+	p.childNodes[name] = pn
+	p.childMu.Unlock()
+}
+
+// removeWithName removes all references with the given name.
 //
-// The provided function is executed after removal.
+// The provided function is executed after reference removal. The only method
+// it may (transitively) call on this pathNode is addChildLocked.
 //
-// Precondition: This call is synchronized w.r.t. other adding or removing of
-// children.
+// If a child pathNode for name exists, it is removed from this pathNode and
+// returned by this function. Any operations on the removed tree must use this
+// value.
 func (p *pathNode) removeWithName(name string, fn func(ref *fidRef)) *pathNode {
-	p.fidRefNames.Range(func(key, value interface{}) bool {
-		if value.(string) == name {
-			p.fidRefNames.Delete(key)
-			fn(key.(*fidRef))
+	p.childMu.Lock()
+	defer p.childMu.Unlock()
+
+	if m, ok := p.childRefs[name]; ok {
+		for ref := range m {
+			delete(m, ref)
+			delete(p.childRefNames, ref)
+			fn(ref)
 		}
-		return true
-	})
+	}
 
-	// Return the original path node.
-	origPathNode := p.pathNodeFor(name)
-	p.childNodes.Delete(name)
+	// Return the original path node, if it exists.
+	origPathNode := p.childNodes[name]
+	delete(p.childNodes, name)
 	return origPathNode
 }