types/lazy: fix flaky TestDeferAfterDo

This test verifies, among other things, that init functions cannot be deferred after (*DeferredFuncs).Do
has already been called and that all subsequent calls to (*DeferredFuncs).Defer return false.

However, the initial implementation of this check was racy: by the time (*DeferredFuncs).Do returned,
not all goroutines that successfully deferred an init function may have incremented the atomic variable
tracking the number of deferred functions. As a result, the variable's value could differ immediately
after (*DeferredFuncs).Do returned and after all goroutines had completed execution (i.e., after wg.Wait()).

In this PR, we replace the original racy check with a different one. Although this new check is also racy,
it can only produce false negatives. This means that if the test fails, it indicates an actual bug rather than
a flaky test.

Fixes #14039

Signed-off-by: Nick Khyl <nickk@tailscale.com>

types/lazy/deferred.go

@@ -22,7 +22,14 @@ type DeferredInit struct {
 // until the owner's [DeferredInit.Do] method is called
 // for the first time.
 //
-// DeferredFuncs is safe for concurrent use.
+// DeferredFuncs is safe for concurrent use. The execution
+// order of functions deferred by different goroutines is
+// unspecified and must not be relied upon.
+// However, functions deferred by the same goroutine are
+// executed in the same relative order they were deferred.
+// Warning: this is the opposite of the behavior of Go's
+// defer statement, which executes deferred functions in
+// reverse order.
 type DeferredFuncs struct {
 	m     sync.Mutex
 	funcs []func() error

types/lazy/deferred_test.go

@@ -205,16 +205,38 @@ func TestDeferredErr(t *testing.T) {
 	}
 }
 
+// TestDeferAfterDo checks all of the following:
+// - Deferring a function before [DeferredInit.Do] is called should always succeed.
+// - All successfully deferred functions are executed by the time [DeferredInit.Do] completes.
+// - No functions can be deferred after [DeferredInit.Do] is called, meaning:
+//   - [DeferredInit.Defer] should return false.
+//   - The deferred function should not be executed.
+//
+// This test is intentionally racy as it attempts to defer functions from multiple goroutines
+// and then calls [DeferredInit.Do] without waiting for them to finish. Waiting would alter
+// the observable behavior and render the test pointless.
 func TestDeferAfterDo(t *testing.T) {
 	var di DeferredInit
 	var deferred, called atomic.Int32
 
+	// deferOnce defers a test function once and fails the test
+	// if [DeferredInit.Defer] returns true after [DeferredInit.Do]
+	// has already been called and any deferred functions have been executed.
+	// It's called concurrently by multiple goroutines.
 	deferOnce := func() bool {
+		// canDefer is whether it's acceptable for Defer to return true.
+		// (but not it necessarily must return true)
+		// If its func has run before, it's definitely not okay for it to
+		// accept more Defer funcs.
+		canDefer := called.Load() == 0
 		ok := di.Defer(func() error {
 			called.Add(1)
 			return nil
 		})
 		if ok {
+			if !canDefer {
+				t.Error("An init function was deferred after DeferredInit.Do() was already called")
+			}
 			deferred.Add(1)
 		}
 		return ok
@@ -242,19 +264,17 @@ func TestDeferAfterDo(t *testing.T) {
 	if err := di.Do(); err != nil {
 		t.Fatalf("DeferredInit.Do() failed: %v", err)
 	}
-	wantDeferred, wantCalled := deferred.Load(), called.Load()
+	// The number of called funcs should remain unchanged after [DeferredInit.Do] returns.
+	wantCalled := called.Load()
 
 	if deferOnce() {
 		t.Error("An init func was deferred after DeferredInit.Do() returned")
 	}
 
 	// Wait for the goroutines deferring init funcs to exit.
-	// No funcs should be deferred after DeferredInit.Do() has returned,
-	// so the deferred and called counters should remain unchanged.
+	// No funcs should be called after DeferredInit.Do() has returned,
+	// and the number of called funcs should be equal to the number of deferred funcs.
 	wg.Wait()
-	if gotDeferred := deferred.Load(); gotDeferred != wantDeferred {
-		t.Errorf("An init func was deferred after DeferredInit.Do() returned. Got %d, want %d", gotDeferred, wantDeferred)
-	}
 	if gotCalled := called.Load(); gotCalled != wantCalled {
 		t.Errorf("An init func was called after DeferredInit.Do() returned. Got %d, want %d", gotCalled, wantCalled)
 	}