runtime: factor waiting on mark phase

There are three places where we wait for the GC mark phase to
complete. Factor these all into a single helper function.

Fixes #24362.

Change-Id: I47f6a7147974f5b9a2869c527a024519070ba6f3
Reviewed-on: https://go-review.googlesource.com/102605
Run-TryBot: Austin Clements <austin@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rick Hudson <rlh@golang.org>

src/runtime/mgc.go

@@ -231,21 +231,10 @@ func setGCPercent(in int32) (out int32) {
 	gcSetTriggerRatio(memstats.triggerRatio)
 	unlock(&mheap_.lock)
 
-	// If we just disabled GC, wait for any concurrent GC to
+	// If we just disabled GC, wait for any concurrent GC mark to
 	// finish so we always return with no GC running.
 	if in < 0 {
-		// Disable phase transitions.
-		lock(&work.sweepWaiters.lock)
-		if gcphase == _GCmark {
-			// GC is active. Wait until we reach sweeping.
-			gp := getg()
-			gp.schedlink = work.sweepWaiters.head
-			work.sweepWaiters.head.set(gp)
-			goparkunlock(&work.sweepWaiters.lock, "wait for GC cycle", traceEvGoBlock, 1)
-		} else {
-			// GC isn't active.
-			unlock(&work.sweepWaiters.lock)
-		}
+		gcWaitOnMark(atomic.Load(&work.cycles))
 	}
 
 	return out
@@ -1090,21 +1079,10 @@ func GC() {
 	// GC may move ahead on its own. For example, when we block
 	// until mark termination N, we may wake up in cycle N+2.
 
-	gp := getg()
-
-	// Prevent the GC phase or cycle count from changing.
-	lock(&work.sweepWaiters.lock)
+	// Wait until the current sweep termination, mark, and mark
+	// termination complete.
 	n := atomic.Load(&work.cycles)
-	if gcphase == _GCmark {
-		// Wait until sweep termination, mark, and mark
-		// termination of cycle N complete.
-		gp.schedlink = work.sweepWaiters.head
-		work.sweepWaiters.head.set(gp)
-		goparkunlock(&work.sweepWaiters.lock, "wait for GC cycle", traceEvGoBlock, 1)
-	} else {
-		// We're in sweep N already.
-		unlock(&work.sweepWaiters.lock)
-	}
+	gcWaitOnMark(n)
 
 	// We're now in sweep N or later. Trigger GC cycle N+1, which
 	// will first finish sweep N if necessary and then enter sweep
@@ -1112,14 +1090,7 @@ func GC() {
 	gcStart(gcBackgroundMode, gcTrigger{kind: gcTriggerCycle, n: n + 1})
 
 	// Wait for mark termination N+1 to complete.
-	lock(&work.sweepWaiters.lock)
-	if gcphase == _GCmark && atomic.Load(&work.cycles) == n+1 {
-		gp.schedlink = work.sweepWaiters.head
-		work.sweepWaiters.head.set(gp)
-		goparkunlock(&work.sweepWaiters.lock, "wait for GC cycle", traceEvGoBlock, 1)
-	} else {
-		unlock(&work.sweepWaiters.lock)
-	}
+	gcWaitOnMark(n + 1)
 
 	// Finish sweep N+1 before returning. We do this both to
 	// complete the cycle and because runtime.GC() is often used
@@ -1156,6 +1127,32 @@ func GC() {
 	releasem(mp)
 }
 
+// gcWaitOnMark blocks until GC finishes the Nth mark phase. If GC has
+// already completed this mark phase, it returns immediately.
+func gcWaitOnMark(n uint32) {
+	for {
+		// Disable phase transitions.
+		lock(&work.sweepWaiters.lock)
+		nMarks := atomic.Load(&work.cycles)
+		if gcphase != _GCmark {
+			// We've already completed this cycle's mark.
+			nMarks++
+		}
+		if nMarks > n {
+			// We're done.
+			unlock(&work.sweepWaiters.lock)
+			return
+		}
+
+		// Wait until sweep termination, mark, and mark
+		// termination of cycle N complete.
+		gp := getg()
+		gp.schedlink = work.sweepWaiters.head
+		work.sweepWaiters.head.set(gp)
+		goparkunlock(&work.sweepWaiters.lock, "wait for GC cycle", traceEvGoBlock, 1)
+	}
+}
+
 // gcMode indicates how concurrent a GC cycle should be.
 type gcMode int