From 8f844ca17bd9b57c217d4757bb5f864b97850915 Mon Sep 17 00:00:00 2001
From: Todd Wright <todd.wright@automattic.com>
Date: Mon, 4 Sep 2023 17:05:29 +1000
Subject: [PATCH] Refactor user stats collector

---
 .../class-user-stats-collector.php            | 91 ++++++++++---------
 1 file changed, 48 insertions(+), 43 deletions(-)

diff --git a/prometheus-collectors/class-user-stats-collector.php b/prometheus-collectors/class-user-stats-collector.php
index 09f38f6923..600bce75d2 100644
--- a/prometheus-collectors/class-user-stats-collector.php
+++ b/prometheus-collectors/class-user-stats-collector.php
@@ -14,10 +14,12 @@ class User_Stats_Collector implements CollectorInterface {
 	private Gauge $role_gauge;
 
 	private string $blog_id;
-	private const TFA_STATUS_ENABLED  = 'enabled';
-	private const TFA_STATUS_DISABLED = 'disabled';
-	private const TFA_STATUS_DEFAULT  = 'unknown';
-	private const METRIC_OPTION       = 'vip-prom-users';
+	private const TFA_STATUS_ENABLED      = 'enabled';
+	private const TFA_STATUS_DISABLED     = 'disabled';
+	private const TFA_STATUS_DEFAULT      = 'unknown';
+	private const METRIC_OPTION           = 'vip-prom-users';
+	private const MAX_USERS_FOR_COUNT     = 1_000_000;
+	private const MAX_USERS_FOR_2FA_STATS = 10_000;
 
 	public function initialize( RegistryInterface $registry ): void {
 		$this->blog_id = Plugin::get_instance()->get_site_label();
@@ -32,19 +34,20 @@ public function initialize( RegistryInterface $registry ): void {
 
 	public function collect_metrics(): void {
 		$metrics = get_option( self::METRIC_OPTION, [] );
-
 		if ( ! $metrics ) {
 			return;
 		}
 
 		foreach ( $metrics as $role => $counts ) {
-			foreach ( $counts['two-factor'] as $status => $count ) {
+			foreach ( $counts as $status => $count ) {
 				$this->role_gauge->set( $count, [ $this->blog_id, $role, $status ] );
 			}
 		}
 	}
 
 	public function process_metrics(): void {
+		global $wpdb;
+
 		if ( ! Context::is_wp_cli() ) {
 			return;
 		}
@@ -54,57 +57,59 @@ public function process_metrics(): void {
 			return;
 		}
 
-		// Limit this to sites without a huge number of users for now
-		if ( wp_is_large_user_count() ) {
+		// count_users() is intensive so skip collecting this metric if there are too many users
+		if ( get_user_count() > self::MAX_USERS_FOR_COUNT ) {
 			return;
 		}
 
-		// Order is not important to us and the query is faster without it
-		add_action( 'pre_user_query', [ $this, 'remove_user_query_orderby' ] );
+		// Get user count by role
+		$user_counts = count_users();
 
-		$users_per_page         = 500;
-		$user_counts            = [];
-		$current_page           = 1;
-		$two_factor_minimum_cap = 'edit_posts';
+		// Get the count of users with 2FA meta to make sure we're not querying too many users with `get_users()`.
+		// This includes users on any network site, not just the current site.
+		// phpcs:ignore WordPress.DB.DirectDatabaseQuery.DirectQuery, WordPress.DB.DirectDatabaseQuery.NoCaching
+		$two_factor_user_count = $wpdb->get_var(
+			$wpdb->prepare( "SELECT COUNT(user_id) FROM {$wpdb->usermeta} WHERE meta_key = %s", \Two_Factor_Core::ENABLED_PROVIDERS_USER_META_KEY )
+		);
 
-		do {
-			$args = [
-				'number'      => $users_per_page,
-				'paged'       => $current_page,
-				'count_total' => false,
+		$include_2fa_stats = $two_factor_user_count < self::MAX_USERS_FOR_2FA_STATS;
+
+		// Prepare the default response
+		foreach ( $user_counts['avail_roles'] as $role => $value ) {
+			$ret[ $role ] = [
+				$include_2fa_stats ? self::TFA_STATUS_DISABLED : self::TFA_STATUS_DEFAULT => $value,
 			];
+		}
 
-			$users = get_users( $args );
+		// If there aren't too many users with the 2FA meta, include 2FA stats
+		if ( $include_2fa_stats ) {
+			// Order is not important to us and the query is faster without the order by clause
+			add_action( 'pre_user_query', [ $this, 'remove_user_query_orderby' ] );
 
-			foreach ( $users as $user ) {
-				$two_factor_status = self::TFA_STATUS_DEFAULT;
+			$users = get_users( [
+				'count_total' => false,
+				'number'      => -1,
+				'meta_key'    => \Two_Factor_Core::ENABLED_PROVIDERS_USER_META_KEY,
+			] );
 
-				// Only get 2fa status for users with the specified cap
-				if ( $user->has_cap( $two_factor_minimum_cap ) ) {
-					$two_factor_status = \Two_Factor_Core::is_user_using_two_factor( $user->ID ) ? self::TFA_STATUS_ENABLED : self::TFA_STATUS_DISABLED;
-				}
+			foreach ( $users as $user ) {
+				// Double-check that 2FA is enabled for this user (the meta could be set to an invalid provider)
+				$two_factor_enabled = \Two_Factor_Core::is_user_using_two_factor( $user->ID );
 
 				foreach ( $user->roles as $role ) {
-					$user_counts[ $role ] ??= [
-						'total'      => 0,
-						'two-factor' => [
-							self::TFA_STATUS_ENABLED  => 0,
-							self::TFA_STATUS_DISABLED => 0,
-							self::TFA_STATUS_DEFAULT  => 0,
-						],
-					];
-					$user_counts[ $role ]['total']++;
-					$user_counts[ $role ]['two-factor'][ $two_factor_status ]++;
+					if ( $two_factor_enabled ) {
+						$ret[ $role ][ self::TFA_STATUS_ENABLED ] ??= 0;
+						$ret[ $role ][ self::TFA_STATUS_ENABLED ]++;
+
+						if ( isset( $ret[ $role ][ self::TFA_STATUS_DISABLED ] ) ) {
+							$ret[ $role ][ self::TFA_STATUS_DISABLED ]--;
+						}
+					}
 				}
 			}
+		}
 
-			$current_page++;
-
-			// Clear the object cache to avoid memory issues
-			vip_reset_local_object_cache();
-		} while ( ! empty( $users ) );
-
-		update_option( self::METRIC_OPTION, $user_counts, false );
+		update_option( self::METRIC_OPTION, $ret, false );
 	}
 
 	public function remove_user_query_orderby( $user_query ) {