fix(menu): restore focus immediately when menu is closed

Currently in some cases we restore focus immediately and in others we wait for the exit animation to finish (when using `matMenuContent`) because the focus restoration logic is coupled to the menu cleanup. This means that if multiple animations overlap, we could end up restoring focus to the wrong element. These changes switch to restoring focus immediately in all cases.

Fixes #16954.

src/material-experimental/mdc-menu/menu.spec.ts

@@ -179,6 +179,25 @@ describe('MatMenu', () => {
     expect(document.activeElement).not.toBe(triggerEl);
   }));
 
+  it('should restore focus to the trigger immediately once the menu is closed', () => {
+    const fixture = createComponent(SimpleMenu, [], [FakeIcon]);
+    fixture.detectChanges();
+    const triggerEl = fixture.componentInstance.triggerEl.nativeElement;
+
+    // A click without a mousedown before it is considered a keyboard open.
+    triggerEl.click();
+    fixture.detectChanges();
+
+    expect(overlayContainerElement.querySelector('.mat-mdc-menu-panel')).toBeTruthy();
+
+    fixture.componentInstance.trigger.closeMenu();
+    fixture.detectChanges();
+    // Note: don't add a `tick` here since we're testing
+    // that focus is restored before the animation is done.
+
+    expect(document.activeElement).toBe(triggerEl);
+  });
+
   it('should be able to set a custom class on the backdrop', fakeAsync(() => {
     const fixture = createComponent(SimpleMenu, [], [FakeIcon]);
 

src/material/menu/menu-trigger.ts

@@ -298,18 +298,20 @@ export class MatMenuTrigger implements AfterContentInit, OnDestroy {
           .subscribe({
             next: () => menu.lazyContent!.detach(),
             // No matter whether the content got re-attached, reset the menu.
-            complete: () => this._resetMenu()
+            complete: () => this._setIsMenuOpen(false)
           });
       } else {
-        this._resetMenu();
+        this._setIsMenuOpen(false);
       }
     } else {
-      this._resetMenu();
+      this._setIsMenuOpen(false);
 
       if (menu.lazyContent) {
         menu.lazyContent.detach();
       }
     }
+
+    this._restoreFocus();
   }
 
   /**
@@ -339,13 +341,8 @@ export class MatMenuTrigger implements AfterContentInit, OnDestroy {
     }
   }
 
-  /**
-   * This method resets the menu when it's closed, most importantly restoring
-   * focus to the menu trigger if the menu was opened via the keyboard.
-   */
-  private _resetMenu(): void {
-    this._setIsMenuOpen(false);
-
+  /** Restores focus to the element that was focused before the menu was open. */
+  private _restoreFocus() {
     // We should reset focus if the user is navigating using a keyboard or
     // if we have a top-level trigger which might cause focus to be lost
     // when clicking on the backdrop.

src/material/menu/menu.spec.ts

@@ -181,6 +181,25 @@ describe('MatMenu', () => {
     expect(document.activeElement).not.toBe(triggerEl);
   }));
 
+  it('should restore focus to the trigger immediately once the menu is closed', () => {
+    const fixture = createComponent(SimpleMenu, [], [FakeIcon]);
+    fixture.detectChanges();
+    const triggerEl = fixture.componentInstance.triggerEl.nativeElement;
+
+    // A click without a mousedown before it is considered a keyboard open.
+    triggerEl.click();
+    fixture.detectChanges();
+
+    expect(overlayContainerElement.querySelector('.mat-menu-panel')).toBeTruthy();
+
+    fixture.componentInstance.trigger.closeMenu();
+    fixture.detectChanges();
+    // Note: don't add a `tick` here since we're testing
+    // that focus is restored before the animation is done.
+
+    expect(document.activeElement).toBe(triggerEl);
+  });
+
   it('should be able to set a custom class on the backdrop', fakeAsync(() => {
     const fixture = createComponent(SimpleMenu, [], [FakeIcon]);