Multisite: Add WP_Site_State class for efficient site state management

src/wp-admin/includes/ms.php

@@ -176,6 +176,8 @@ function wpmu_delete_user( $id ) {
 	$blogs = get_blogs_of_user( $id );
 
 	if ( ! empty( $blogs ) ) {
+		$original_state = get_site_state();
+
 		foreach ( $blogs as $blog ) {
 			switch_to_blog( $blog->userblog_id );
 			remove_user_from_blog( $id, $blog->userblog_id );
@@ -193,9 +195,10 @@ function wpmu_delete_user( $id ) {
 					wp_delete_link( $link_id );
 				}
 			}
-
-			restore_current_blog();
 		}
+
+		// Restore the original site state once after processing all sites.
+		restore_site_state( $original_state );
 	}
 
 	$meta = $wpdb->get_col( $wpdb->prepare( "SELECT umeta_id FROM $wpdb->usermeta WHERE user_id = %d", $id ) );

src/wp-includes/admin-bar.php

@@ -672,6 +672,9 @@ function wp_admin_bar_my_sites_menu( $wp_admin_bar ) {
 	 */
 	$show_site_icons = apply_filters( 'wp_admin_bar_show_site_icons', true );
 
+	// Store the current site state before iterating through user blogs.
+	$original_state = get_site_state();
+
 	foreach ( (array) $wp_admin_bar->user->blogs as $blog ) {
 		switch_to_blog( $blog->userblog_id );
 
@@ -753,9 +756,10 @@ function wp_admin_bar_my_sites_menu( $wp_admin_bar ) {
 				'href'   => home_url( '/' ),
 			)
 		);
-
-		restore_current_blog();
 	}
+
+	// Restore the original site state once after processing all blogs.
+	restore_site_state( $original_state );
 }
 
 /**

src/wp-includes/class-wp-site-state.php

@@ -0,0 +1,139 @@
+<?php
+/**
+ * Site State API: WP_Site_State class
+ *
+ * @package WordPress
+ * @subpackage Multisite
+ * @since 6.8.0
+ */
+
+/**
+ * Core class used for efficiently switching and restoring site state.
+ *
+ * @since 6.8.0
+ */
+#[AllowDynamicProperties]
+class WP_Site_State {
+	/**
+	 * Current site ID.
+	 *
+	 * @since 6.8.0
+	 * @var int
+	 */
+	private $site_id;
+
+	/**
+	 * The switched stack.
+	 *
+	 * @since 6.8.0
+	 * @var array
+	 */
+	private $switched_stack = array();
+
+	/**
+	 * Whether or not we're currently switched.
+	 *
+	 * @since 6.8.0
+	 * @var bool
+	 */
+	private $switched = false;
+
+	/**
+	 * Constructor.
+	 *
+	 * Stores the current site ID, the switched stack, and the switched state.
+	 *
+	 * @since 6.8.0
+	 */
+	public function __construct() {
+		global $_wp_switched_stack, $switched;
+
+		$this->site_id = get_current_blog_id();
+
+		if ( ! empty( $_wp_switched_stack ) ) {
+			$this->switched_stack = $_wp_switched_stack;
+		}
+
+		$this->switched = ! empty( $switched );
+	}
+
+	/**
+	 * Restores the stored site state.
+	 *
+	 * @since 6.8.0
+	 *
+	 * @return bool True on success, false if no state change was needed.
+	 */
+	public function restore() {
+		global $_wp_switched_stack, $switched, $wpdb, $blog_id, $table_prefix;
+
+		$current_blog_id = get_current_blog_id();
+
+		// If we're already on the target blog, just update the global state.
+		if ( $current_blog_id === $this->site_id ) {
+			$_wp_switched_stack = $this->switched_stack;
+			$switched           = $this->switched;
+			return true;
+		}
+
+		$wpdb->set_blog_id( $this->site_id );
+		$table_prefix = $wpdb->get_blog_prefix();
+		$blog_id      = $this->site_id;
+
+		if ( function_exists( 'wp_cache_switch_to_blog' ) ) {
+			wp_cache_switch_to_blog( $blog_id );
+		}
+
+		// Restore the switched stack and state.
+		$_wp_switched_stack = $this->switched_stack;
+		$switched           = $this->switched;
+
+		/**
+		 * Fires when the blog is switched.
+		 *
+		 * @since MU (3.0.0)
+		 * @since 5.4.0 The `$context` parameter was added.
+		 *
+		 * @param int    $new_blog_id  New blog ID.
+		 * @param int    $prev_blog_id Previous blog ID.
+		 * @param string $context      Additional context. Accepts 'switch' when called from switch_to_blog()
+		 *                             or 'restore' when called from restore_current_blog().
+		 */
+		do_action( 'switch_blog', $blog_id, $current_blog_id, 'restore_state' );
+
+		return true;
+	}
+
+	/**
+	 * Gets the site ID stored in this state.
+	 *
+	 * @since 6.8.0
+	 *
+	 * @return int The site ID.
+	 */
+	public function get_site_id() {
+		return $this->site_id;
+	}
+
+	/**
+	 * Gets the switched stack stored in this state.
+	 *
+	 * @since 6.8.0
+	 *
+	 * @return array The switched stack.
+	 */
+	public function get_switched_stack() {
+		return $this->switched_stack;
+	}
+
+	/**
+	 * Gets the switched status stored in this state.
+	 *
+	 * @since 6.8.0
+	 *
+	 * @return bool Whether the site was switched.
+	 */
+	public function is_switched() {
+		return $this->switched;
+	}
+}

src/wp-includes/class-wp-xmlrpc-server.php

@@ -753,6 +753,7 @@ public function wp_getUsersBlogs( $args ) {
 		}
 
 		$current_network_id = get_current_network_id();
+		$original_state = get_site_state();
 
 		foreach ( $blogs as $blog ) {
 			// Don't include blogs that aren't hosted at this site.
@@ -775,10 +776,11 @@ public function wp_getUsersBlogs( $args ) {
 				'blogName'  => get_option( 'blogname' ),
 				'xmlrpc'    => site_url( 'xmlrpc.php', 'rpc' ),
 			);
-
-			restore_current_blog();
 		}
 
+		// Restore the original site state once after processing all blogs.
+		restore_site_state( $original_state );
+
 		return $struct;
 	}
 

src/wp-includes/ms-blogs.php

@@ -970,3 +970,37 @@ function wp_count_sites( $network_id = null ) {
 
 	return $counts;
 }
+
+/**
+ * Creates a snapshot of the current site state.
+ *
+ * This function creates a WP_Site_State object that captures the current site
+ * state, including the site ID, switched stack, and switched status.
+ *
+ * @since 6.8.0
+ *
+ * @return WP_Site_State A snapshot of the current site state.
+ */
+function get_site_state() {
+	return new WP_Site_State();
+}
+
+/**
+ * Restores a previously saved site state.
+ *
+ * This function efficiently restores a site state that was captured using
+ * get_site_state(), without unnecessary intermediate restores when switching
+ * between multiple sites.
+ *
+ * @since 6.8.0
+ *
+ * @param WP_Site_State $state The site state object to restore to.
+ * @return bool True on success, false on failure.
+ */
+function restore_site_state( $state ) {
+	if ( ! $state instanceof WP_Site_State ) {
+		return false;
+	}
+
+	return $state->restore();
+}

src/wp-includes/ms-functions.php

@@ -138,9 +138,12 @@ function get_blog_count( $network_id = null ) {
  * @return WP_Post|null WP_Post object on success, null on failure
  */
 function get_blog_post( $blog_id, $post_id ) {
+	$original_state = get_site_state();
+
 	switch_to_blog( $blog_id );
 	$post = get_post( $post_id );
-	restore_current_blog();
+
+	restore_site_state( $original_state );
 
 	return $post;
 }

src/wp-includes/ms-settings.php

@@ -42,6 +42,9 @@
 /** WP_Site class */
 require_once ABSPATH . WPINC . '/class-wp-site.php';
 
+/** WP_Site_State class */
+require_once ABSPATH . WPINC . '/class-wp-site-state.php';
+
 /** Multisite loader */
 require_once ABSPATH . WPINC . '/ms-load.php';
 

tests/phpunit/tests/multisite/wpSiteState.php

@@ -0,0 +1,94 @@
+<?php
+
+if ( is_multisite() ) :
+
+	/**
+	 * @group multisite
+	 * @group ms-site
+	 */
+	class Tests_Multisite_WpSiteState extends WP_UnitTestCase {
+		protected static $site_ids;
+		protected static $network_id;
+
+		public static function wpSetUpBeforeClass( WP_UnitTest_Factory $factory ) {
+			self::$network_id = $factory->network->create();
+			self::$site_ids   = array();
+
+			for ( $i = 0; $i < 3; $i++ ) {
+				self::$site_ids[] = $factory->blog->create_object(
+					array(
+						'domain'  => 'wordpress.org',
+						'path'    => '/sites/' . $i,
+						'site_id' => self::$network_id,
+					)
+				);
+			}
+		}
+
+		/**
+		 * Tests that the site state can be saved and restored.
+		 *
+		 * @ticket 37958
+		 */
+		public function test_site_state_save_and_restore() {
+			$original_blog_id = get_current_blog_id();
+
+			$state = get_site_state();
+
+			switch_to_blog( self::$site_ids[0] );
+			$this->assertEquals( self::$site_ids[0], get_current_blog_id() );
+
+			switch_to_blog( self::$site_ids[1] );
+			$this->assertEquals( self::$site_ids[1], get_current_blog_id() );
+
+			restore_site_state( $state );
+			$this->assertEquals( $original_blog_id, get_current_blog_id() );
+
+			$this->assertFalse( ms_is_switched() );
+		}
+
+		/**
+		 * Tests that the site state can be saved and restored in bulk operations.
+		 *
+		 * @ticket 37958
+		 */
+		public function test_site_state_in_bulk_operations() {
+			$original_blog_id = get_current_blog_id();
+
+			$state = get_site_state();
+
+			foreach ( self::$site_ids as $site_id ) {
+				switch_to_blog( $site_id );
+				$this->assertEquals( $site_id, get_current_blog_id() );
+			}
+
+			restore_site_state( $state );
+			$this->assertEquals( $original_blog_id, get_current_blog_id() );
+
+			$this->assertFalse( ms_is_switched() );
+		}
+
+		/**
+		 * Tests that the site state object maintains its properties.
+		 *
+		 * @ticket 37958
+		 */
+		public function test_site_state_maintains_properties() {
+			$state = get_site_state();
+
+			$this->assertEquals( get_current_blog_id(), $state->get_site_id() );
+			$this->assertEquals( ms_is_switched(), $state->is_switched() );
+
+			switch_to_blog( self::$site_ids[0] );
+			$switched_state = get_site_state();
+			$this->assertTrue( $switched_state->is_switched() );
+
+			restore_site_state( $switched_state );
+			$this->assertEquals( self::$site_ids[0], get_current_blog_id() );
+
+			restore_site_state( $state );
+			$this->assertFalse( ms_is_switched() );
+		}
+	}
+
+endif;