runtime: merge all treaps into one implementation

This change modifies the treap implementation to support holding all
spans in a single treap, instead of keeping them all in separate treaps.

This improves ergonomics for nearly all treap-related callsites.
With that said, iteration is now more expensive, but it never occurs on
the fast path, only on scavenging-related paths.

This change opens up the opportunity for further optimizations, such as
splitting spans without treap removal (taking treap removal off the span
allocator's critical path) as well as improvements to treap iteration
(building linked lists for each iteration type and managing them on
insert/removal, since those operations should be less frequent).

For #30333.

Change-Id: I3dac97afd3682a37fda09ae8656a770e1369d0a9
Reviewed-on: https://go-review.googlesource.com/c/go/+/174398
Run-TryBot: Michael Knyszek <mknyszek@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Austin Clements <austin@google.com>

src/runtime/export_test.go

@@ -339,7 +339,7 @@ func ReadMemStatsSlow() (base, slow MemStats) {
 			slow.BySize[i].Frees = bySize[i].Frees
 		}
 
-		for i := mheap_.scav.start(); i.valid(); i = i.next() {
+		for i := mheap_.free.start(0, 0); i.valid(); i = i.next() {
 			slow.HeapReleased += uint64(i.span().released())
 		}
 
@@ -522,11 +522,12 @@ type Span struct {
 	*mspan
 }
 
-func AllocSpan(base, npages uintptr) Span {
+func AllocSpan(base, npages uintptr, scavenged bool) Span {
 	lock(&mheap_.lock)
 	s := (*mspan)(mheap_.spanalloc.alloc())
 	unlock(&mheap_.lock)
 	s.init(base, npages)
+	s.scavenged = scavenged
 	return Span{s}
 }
 
@@ -545,6 +546,17 @@ func (s Span) Pages() uintptr {
 	return s.mspan.npages
 }
 
+type TreapIterType int
+
+const (
+	TreapIterScav TreapIterType = TreapIterType(treapIterScav)
+	TreapIterBits               = treapIterBits
+)
+
+func (s Span) MatchesIter(mask, match TreapIterType) bool {
+	return s.mspan.matchesIter(treapIterType(mask), treapIterType(match))
+}
+
 type TreapIter struct {
 	treapIter
 }
@@ -575,12 +587,12 @@ type Treap struct {
 	mTreap
 }
 
-func (t *Treap) Start() TreapIter {
-	return TreapIter{t.start()}
+func (t *Treap) Start(mask, match TreapIterType) TreapIter {
+	return TreapIter{t.start(treapIterType(mask), treapIterType(match))}
 }
 
-func (t *Treap) End() TreapIter {
-	return TreapIter{t.end()}
+func (t *Treap) End(mask, match TreapIterType) TreapIter {
+	return TreapIter{t.end(treapIterType(mask), treapIterType(match))}
 }
 
 func (t *Treap) Insert(s Span) {

src/runtime/mgclarge.go

@@ -185,6 +185,41 @@ func (t *treapNode) validateMaxPages() uintptr {
 	return max
 }
 
+// treapIterType represents the type of iteration to perform
+// over the treap. Each choice effectively represents a filter,
+// i.e. spans that do not satisfy the conditions of the iteration
+// type will be skipped over.
+type treapIterType uint8
+
+const (
+	treapIterScav treapIterType = 1 << iota // scavenged spans
+	treapIterBits               = iota
+)
+
+// matches returns true if t satisfies the filter given by mask and match. mask
+// is a bit-set of span properties to filter on.
+//
+// In other words, matches returns true if all properties set in mask have the
+// value given by the corresponding bits in match.
+func (t treapIterType) matches(mask, match treapIterType) bool {
+	return t&mask == match
+}
+
+// iterType returns the treapIterType associated with this span.
+func (s *mspan) iterType() treapIterType {
+	have := treapIterType(0)
+	if s.scavenged {
+		have |= treapIterScav
+	}
+	return have
+}
+
+// matchesIter is a convenience method which checks if a span
+// meets the criteria of the mask and match for an iter type.
+func (s *mspan) matchesIter(mask, match treapIterType) bool {
+	return s.iterType().matches(mask, match)
+}
+
 // treapIter is a bidirectional iterator type which may be used to iterate over a
 // an mTreap in-order forwards (increasing order) or backwards (decreasing order).
 // Its purpose is to hide details about the treap from users when trying to iterate
@@ -192,6 +227,7 @@ func (t *treapNode) validateMaxPages() uintptr {
 //
 // To create iterators over the treap, call start or end on an mTreap.
 type treapIter struct {
+	mask, match treapIterType
 	t           *treapNode
 }
 
@@ -211,6 +247,9 @@ func (i *treapIter) valid() bool {
 // ceases to be valid, calling next will panic.
 func (i treapIter) next() treapIter {
 	i.t = i.t.succ()
+	for i.valid() && !i.span().matchesIter(i.mask, i.match) {
+		i.t = i.t.succ()
+	}
 	return i
 }
 
@@ -218,12 +257,15 @@ func (i treapIter) next() treapIter {
 // ceases to be valid, calling prev will panic.
 func (i treapIter) prev() treapIter {
 	i.t = i.t.pred()
+	for i.valid() && !i.span().matchesIter(i.mask, i.match) {
+		i.t = i.t.pred()
+	}
 	return i
 }
 
 // start returns an iterator which points to the start of the treap (the
-// left-most node in the treap).
-func (root *mTreap) start() treapIter {
+// left-most node in the treap) subject to mask and match constraints.
+func (root *mTreap) start(mask, match treapIterType) treapIter {
 	t := root.treap
 	if t == nil {
 		return treapIter{}
@@ -231,12 +273,15 @@ func (root *mTreap) start() treapIter {
 	for t.left != nil {
 		t = t.left
 	}
-	return treapIter{t: t}
+	for t != nil && !t.span.matchesIter(mask, match) {
+		t = t.succ()
+	}
+	return treapIter{mask, match, t}
 }
 
 // end returns an iterator which points to the end of the treap (the
-// right-most node in the treap).
-func (root *mTreap) end() treapIter {
+// right-most node in the treap) subject to mask and match constraints.
+func (root *mTreap) end(mask, match treapIterType) treapIter {
 	t := root.treap
 	if t == nil {
 		return treapIter{}
@@ -244,7 +289,10 @@ func (root *mTreap) end() treapIter {
 	for t.right != nil {
 		t = t.right
 	}
-	return treapIter{t: t}
+	for t != nil && !t.span.matchesIter(mask, match) {
+		t = t.pred()
+	}
+	return treapIter{mask, match, t}
 }
 
 // insert adds span to the large span treap.
@@ -348,9 +396,10 @@ func (root *mTreap) removeNode(t *treapNode) {
 	mheap_.treapalloc.free(unsafe.Pointer(t))
 }
 
-// find searches for, finds, and returns the treap iterator representing the
-// position of the span with the smallest base address which is at least npages
-// in size. If no span has at least npages it returns an invalid iterator.
+// find searches for, finds, and returns the treap iterator over all spans
+// representing the position of the span with the smallest base address which is
+// at least npages in size. If no span has at least npages it returns an invalid
+// iterator.
 //
 // This algorithm is as follows:
 // * If there's a left child and its subtree can satisfy this allocation,
@@ -391,7 +440,7 @@ func (root *mTreap) find(npages uintptr) treapIter {
 			t = nil
 		}
 	}
-	return treapIter{t}
+	return treapIter{t: t}
 }
 
 // removeSpan searches for, finds, deletes span along with

src/runtime/mheap.go

@@ -30,8 +30,7 @@ const minPhysPageSize = 4096
 //go:notinheap
 type mheap struct {
 	lock      mutex
-	free      mTreap // free and non-scavenged spans
-	scav      mTreap // free and scavenged spans
+	free      mTreap // free spans
 	sweepgen  uint32 // sweep generation, see comment in mspan
 	sweepdone uint32 // all spans are swept
 	sweepers  uint32 // number of active sweepone calls
@@ -60,7 +59,7 @@ type mheap struct {
 	// on the swept stack.
 	sweepSpans [2]gcSweepBuf
 
-	_ uint32 // align uint64 fields on 32-bit for atomics
+	// _ uint32 // align uint64 fields on 32-bit for atomics
 
 	// Proportional sweep
 	//
@@ -464,7 +463,7 @@ func (h *mheap) coalesce(s *mspan) {
 
 		// The size is potentially changing so the treap needs to delete adjacent nodes and
 		// insert back as a combined node.
-		h.treapForSpan(other).removeSpan(other)
+		h.free.removeSpan(other)
 		other.state = mSpanDead
 		h.spanalloc.free(unsafe.Pointer(other))
 	}
@@ -482,7 +481,7 @@ func (h *mheap) coalesce(s *mspan) {
 			return
 		}
 		// Since we're resizing other, we must remove it from the treap.
-		h.treapForSpan(other).removeSpan(other)
+		h.free.removeSpan(other)
 
 		// Round boundary to the nearest physical page size, toward the
 		// scavenged span.
@@ -500,7 +499,7 @@ func (h *mheap) coalesce(s *mspan) {
 		h.setSpan(boundary, b)
 
 		// Re-insert other now that it has a new size.
-		h.treapForSpan(other).insert(other)
+		h.free.insert(other)
 	}
 
 	// Coalesce with earlier, later spans.
@@ -1101,57 +1100,27 @@ func (h *mheap) setSpans(base, npage uintptr, s *mspan) {
 	}
 }
 
-// treapForSpan returns the appropriate treap for a span for
-// insertion and removal.
-func (h *mheap) treapForSpan(span *mspan) *mTreap {
-	if span.scavenged {
-		return &h.scav
-	}
-	return &h.free
-}
-
-// pickFreeSpan acquires a free span from internal free list
-// structures if one is available. Otherwise returns nil.
-// h must be locked.
-func (h *mheap) pickFreeSpan(npage uintptr) *mspan {
-	tf := h.free.find(npage)
-	ts := h.scav.find(npage)
-
-	// Check for whichever treap gave us the smaller, non-nil result.
-	// Note that we want the _smaller_ free span, i.e. the free span
-	// closer in size to the amount we requested (npage).
-	var s *mspan
-	if tf.valid() && (!ts.valid() || tf.span().base() <= ts.span().base()) {
-		s = tf.span()
-		h.free.erase(tf)
-	} else if ts.valid() && (!tf.valid() || tf.span().base() > ts.span().base()) {
-		s = ts.span()
-		h.scav.erase(ts)
-	}
-	return s
-}
-
 // Allocates a span of the given size.  h must be locked.
 // The returned span has been removed from the
 // free structures, but its state is still mSpanFree.
 func (h *mheap) allocSpanLocked(npage uintptr, stat *uint64) *mspan {
-	var s *mspan
-
-	s = h.pickFreeSpan(npage)
-	if s != nil {
+	t := h.free.find(npage)
+	if t.valid() {
 		goto HaveSpan
 	}
-	// On failure, grow the heap and try again.
 	if !h.grow(npage) {
 		return nil
 	}
-	s = h.pickFreeSpan(npage)
-	if s != nil {
+	t = h.free.find(npage)
+	if t.valid() {
 		goto HaveSpan
 	}
 	throw("grew heap, but no adequate free span found")
 
 HaveSpan:
+	s := t.span()
+	h.free.erase(t)
+
 	// Mark span in use.
 	if s.state != mSpanFree {
 		throw("candidate mspan for allocation is not free")
@@ -1339,8 +1308,8 @@ func (h *mheap) freeSpanLocked(s *mspan, acctinuse, acctidle bool, unusedsince i
 	// Coalesce span with neighbors.
 	h.coalesce(s)
 
-	// Insert s into the appropriate treap.
-	h.treapForSpan(s).insert(s)
+	// Insert s into the treap.
+	h.free.insert(s)
 }
 
 // scavengeLocked scavenges nbytes worth of spans in the free treap by
@@ -1355,10 +1324,11 @@ func (h *mheap) scavengeLocked(nbytes uintptr) {
 		h.scavengeCredit -= nbytes
 		return
 	}
-	// Iterate over the treap backwards (from highest address to lowest address)
-	// scavenging spans until we've reached our quota of nbytes.
+	// Iterate over the unscavenged spans in the treap backwards (from highest
+	// address to lowest address) scavenging spans until we've reached our
+	// quota of nbytes.
 	released := uintptr(0)
-	for t := h.free.end(); released < nbytes && t.valid(); {
+	for t := h.free.end(treapIterScav, 0); released < nbytes && t.valid(); {
 		s := t.span()
 		r := s.scavenge()
 		if r == 0 {
@@ -1373,7 +1343,7 @@ func (h *mheap) scavengeLocked(nbytes uintptr) {
 		// the same scavenged state adjacent to each other.
 		h.coalesce(s)
 		t = n
-		h.scav.insert(s)
+		h.free.insert(s)
 		released += r
 	}
 	// If we over-scavenged, turn that extra amount into credit.
@@ -1386,20 +1356,19 @@ func (h *mheap) scavengeLocked(nbytes uintptr) {
 // treapNode's span. It then removes the scavenged span from
 // unscav and adds it into scav before continuing. h must be locked.
 func (h *mheap) scavengeAllLocked(now, limit uint64) uintptr {
-	// Iterate over the treap scavenging spans if unused for at least limit time.
+	// Iterate over the unscavenged spans in the treap scavenging spans
+	// if unused for at least limit time.
 	released := uintptr(0)
-	for t := h.free.start(); t.valid(); {
+	for t := h.free.start(treapIterScav, 0); t.valid(); {
 		s := t.span()
 		n := t.next()
 		if (now - uint64(s.unusedsince)) > limit {
 			r := s.scavenge()
 			if r != 0 {
 				h.free.erase(t)
-				// Now that s is scavenged, we must eagerly coalesce it
-				// with its neighbors to prevent having two spans with
-				// the same scavenged state adjacent to each other.
+				// See (*mheap).scavenge.
 				h.coalesce(s)
-				h.scav.insert(s)
+				h.free.insert(s)
 				released += r
 			}
 		}

src/runtime/treap_test.go

@@ -5,19 +5,23 @@
 package runtime_test
 
 import (
+	"fmt"
 	"runtime"
 	"testing"
 )
 
-var spanDesc = map[uintptr]uintptr{
-	0xc0000000: 2,
-	0xc0006000: 1,
-	0xc0010000: 8,
-	0xc0022000: 7,
-	0xc0034000: 4,
-	0xc0040000: 5,
-	0xc0050000: 5,
-	0xc0060000: 5000,
+var spanDesc = map[uintptr]struct {
+	pages uintptr
+	scav  bool
+}{
+	0xc0000000: {2, false},
+	0xc0006000: {1, false},
+	0xc0010000: {8, false},
+	0xc0022000: {7, false},
+	0xc0034000: {4, true},
+	0xc0040000: {5, false},
+	0xc0050000: {5, true},
+	0xc0060000: {5000, false},
 }
 
 // Wrap the Treap one more time because go:notinheap doesn't
@@ -28,6 +32,10 @@ type treap struct {
 	runtime.Treap
 }
 
+func maskMatchName(mask, match runtime.TreapIterType) string {
+	return fmt.Sprintf("%0*b-%0*b", runtime.TreapIterBits, uint8(mask), runtime.TreapIterBits, uint8(match))
+}
+
 // This test ensures that the treap implementation in the runtime
 // maintains all stated invariants after different sequences of
 // insert, removeSpan, find, and erase. Invariants specific to the
@@ -36,12 +44,37 @@ type treap struct {
 // treap.
 func TestTreap(t *testing.T) {
 	// Set up a bunch of spans allocated into mheap_.
+	// Also, derive a set of typeCounts of each type of span
+	// according to runtime.TreapIterType so we can verify against
+	// them later.
 	spans := make([]runtime.Span, 0, len(spanDesc))
-	for base, pages := range spanDesc {
-		s := runtime.AllocSpan(base, pages)
+	typeCounts := [1 << runtime.TreapIterBits][1 << runtime.TreapIterBits]int{}
+	for base, de := range spanDesc {
+		s := runtime.AllocSpan(base, de.pages, de.scav)
 		defer s.Free()
 		spans = append(spans, s)
+
+		for i := runtime.TreapIterType(0); i < 1<<runtime.TreapIterBits; i++ {
+			for j := runtime.TreapIterType(0); j < 1<<runtime.TreapIterBits; j++ {
+				if s.MatchesIter(i, j) {
+					typeCounts[i][j]++
+				}
+			}
+		}
+	}
+	t.Run("TypeCountsSanity", func(t *testing.T) {
+		// Just sanity check type counts for a few values.
+		check := func(mask, match runtime.TreapIterType, count int) {
+			tc := typeCounts[mask][match]
+			if tc != count {
+				name := maskMatchName(mask, match)
+				t.Fatalf("failed a sanity check for mask/match %s counts: got %d, wanted %d", name, tc, count)
+			}
 		}
+		check(0, 0, len(spanDesc))
+		check(runtime.TreapIterScav, 0, 6)
+		check(runtime.TreapIterScav, runtime.TreapIterScav, 2)
+	})
 	t.Run("Insert", func(t *testing.T) {
 		tr := treap{}
 		// Test just a very basic insert/remove for sanity.
@@ -77,24 +110,33 @@ func TestTreap(t *testing.T) {
 		}
 	})
 	t.Run("Iterate", func(t *testing.T) {
+		for mask := runtime.TreapIterType(0); mask < 1<<runtime.TreapIterBits; mask++ {
+			for match := runtime.TreapIterType(0); match < 1<<runtime.TreapIterBits; match++ {
+				iterName := maskMatchName(mask, match)
+				t.Run(iterName, func(t *testing.T) {
 					t.Run("StartToEnd", func(t *testing.T) {
 						// Ensure progressing an iterator actually goes over the whole treap
 						// from the start and that it iterates over the elements in order.
-			// Also ensures that Start returns a valid iterator.
+						// Furthermore, ensure that it only iterates over the relevant parts
+						// of the treap.
+						// Finally, ensures that Start returns a valid iterator.
 						tr := treap{}
 						for _, s := range spans {
 							tr.Insert(s)
 						}
 						nspans := 0
 						lastBase := uintptr(0)
-			for i := tr.Start(); i.Valid(); i = i.Next() {
+						for i := tr.Start(mask, match); i.Valid(); i = i.Next() {
 							nspans++
 							if lastBase > i.Span().Base() {
 								t.Fatalf("not iterating in correct order: encountered base %x before %x", lastBase, i.Span().Base())
 							}
 							lastBase = i.Span().Base()
+							if !i.Span().MatchesIter(mask, match) {
+								t.Fatalf("found non-matching span while iteration over mask/match %s: base %x", iterName, i.Span().Base())
+							}
 						}
-			if nspans != len(spans) {
+						if nspans != typeCounts[mask][match] {
 							t.Fatal("failed to iterate forwards over full treap")
 						}
 						for _, s := range spans {
@@ -102,36 +144,40 @@ func TestTreap(t *testing.T) {
 						}
 					})
 					t.Run("EndToStart", func(t *testing.T) {
-			// Ensure progressing an iterator actually goes over the whole treap
-			// from the end and that it iterates over the elements in reverse
-			// order. Also ensures that End returns a valid iterator.
+						// See StartToEnd tests.
 						tr := treap{}
 						for _, s := range spans {
 							tr.Insert(s)
 						}
 						nspans := 0
 						lastBase := ^uintptr(0)
-			for i := tr.End(); i.Valid(); i = i.Prev() {
+						for i := tr.End(mask, match); i.Valid(); i = i.Prev() {
 							nspans++
 							if lastBase < i.Span().Base() {
 								t.Fatalf("not iterating in correct order: encountered base %x before %x", lastBase, i.Span().Base())
 							}
 							lastBase = i.Span().Base()
+							if !i.Span().MatchesIter(mask, match) {
+								t.Fatalf("found non-matching span while iteration over mask/match %s: base %x", iterName, i.Span().Base())
+							}
 						}
-			if nspans != len(spans) {
+						if nspans != typeCounts[mask][match] {
 							t.Fatal("failed to iterate backwards over full treap")
 						}
 						for _, s := range spans {
 							tr.RemoveSpan(s)
 						}
 					})
+				})
+			}
+		}
 		t.Run("Prev", func(t *testing.T) {
 			// Test the iterator invariant that i.prev().next() == i.
 			tr := treap{}
 			for _, s := range spans {
 				tr.Insert(s)
 			}
-			i := tr.Start().Next().Next()
+			i := tr.Start(0, 0).Next().Next()
 			p := i.Prev()
 			if !p.Valid() {
 				t.Fatal("i.prev() is invalid")
@@ -149,7 +195,7 @@ func TestTreap(t *testing.T) {
 			for _, s := range spans {
 				tr.Insert(s)
 			}
-			i := tr.Start().Next().Next()
+			i := tr.Start(0, 0).Next().Next()
 			n := i.Next()
 			if !n.Valid() {
 				t.Fatal("i.next() is invalid")
@@ -169,7 +215,7 @@ func TestTreap(t *testing.T) {
 		for _, s := range spans {
 			tr.Insert(s)
 		}
-		i := tr.Start().Next().Next().Next()
+		i := tr.Start(0, 0).Next().Next().Next()
 		s := i.Span()
 		n := i.Next()
 		p := i.Prev()
@@ -191,7 +237,7 @@ func TestTreap(t *testing.T) {
 		for _, s := range spans {
 			tr.Insert(s)
 		}
-		for i := tr.Start(); i.Valid(); {
+		for i := tr.Start(0, 0); i.Valid(); {
 			n := i.Next()
 			tr.Erase(i)
 			i = n