Look at a module's mark again after letting go of it, so a wake just before is taken (issue 1945)

cmd/mesh-controller/composed_keys.go

@@ -120,6 +120,10 @@ var (
 	doneKeysSay     = func(s string) { fmt.Println(s) }
 )
 
+// doneKeysBeforeLetGo is called just before a follower lets go of its module; a variable so a test can wake the module
+// at that moment. Nil does nothing.
+var doneKeysBeforeLetGo func(module string)
+
 // doneKeysLookAgain marks a module woken while a take-in of it ran: that take-in goes round once more.
 var doneKeysLookAgain sync.Map
 
@@ -132,13 +136,29 @@ var doneKeysMu sync.Map
 // yet is no failure: it is said once for each reason it waits for and looked at again every doneKeyRetry.
 func takeDoneKey(ctx context.Context, open *stores, module string) {
 	mu, _ := doneKeysMu.LoadOrStore(module, &sync.Mutex{})
-	if !mu.(*sync.Mutex).TryLock() {
-		// A take-in of the module already follows it: it is told to look again once it is done, so a newer key written
-		// meanwhile is never lost (the review of issue 1945).
-		doneKeysLookAgain.Store(module, true)
-		return
+	for {
+		if !mu.(*sync.Mutex).TryLock() {
+			// A take-in of the module already follows it: it is told to look again once it is done, so a newer key
+			// written meanwhile is never lost (the review of issue 1945).
+			doneKeysLookAgain.Store(module, true)
+			return
+		}
+		followDoneKey(ctx, open, module)
+		if doneKeysBeforeLetGo != nil {
+			doneKeysBeforeLetGo(module)
+		}
+		mu.(*sync.Mutex).Unlock()
+		// A wake between this follower's last look and its letting go failed its TryLock and only marked the module:
+		// looked at again now that it is let go, and taken once more (the review of issue 1945).
+		if _, again := doneKeysLookAgain.LoadAndDelete(module); !again || ctx.Err() != nil {
+			return
+		}
 	}
-	defer mu.(*sync.Mutex).Unlock()
+}
+
+// followDoneKey takes a module's done key in, under the module's lock, until it is taken in or ctx ends: a failure or a
+// wait is tried again every doneKeyRetry.
+func followDoneKey(ctx context.Context, open *stores, module string) {
 	failed := false
 	saidWait := ""
 	for {

cmd/mesh-controller/composed_keys_waits_test.go

@@ -121,3 +121,35 @@ func TestAWakeDuringATakeInIsTakenAfterIt(t *testing.T) {
 		t.Fatalf("taken %d times: the wake during the take-in was lost", calls)
 	}
 }
+
+// A wake that comes after the follower's last look but before it lets go of the module is not lost: the follower looks
+// at the mark again once it has let go, and takes the module once more (the review of issue 1945).
+func TestAWakeJustBeforeTheFollowerLetsGoIsTaken(t *testing.T) {
+	was, wasRetry, wasSay, wasHook := doneKeyTakeOnce, doneKeyRetry, doneKeysSay, doneKeysBeforeLetGo
+	t.Cleanup(func() {
+		doneKeyTakeOnce, doneKeyRetry, doneKeysSay, doneKeysBeforeLetGo = was, wasRetry, wasSay, wasHook
+	})
+	doneKeyRetry = 10 * time.Millisecond
+	doneKeysSay = func(string) {}
+	var mu sync.Mutex
+	calls := 0
+	doneKeyTakeOnce = func(context.Context, *stores, string) error {
+		mu.Lock()
+		calls++
+		mu.Unlock()
+		return nil
+	}
+	woke := false
+	doneKeysBeforeLetGo = func(module string) {
+		if !woke {
+			woke = true
+			takeDoneKey(context.Background(), nil, module) // fails its TryLock and marks the module
+		}
+	}
+	takeDoneKey(context.Background(), nil, "mesh-host")
+	mu.Lock()
+	defer mu.Unlock()
+	if calls != 2 {
+		t.Fatalf("taken %d times: the wake just before the follower let go was lost", calls)
+	}
+}