Improved active user count fetching for EE token checks (#29400)

* Improved active user count fetching

* Update tests.

src/metabase/public_settings/premium_features.clj

@@ -16,6 +16,7 @@
    [metabase.util.log :as log]
    [metabase.util.schema :as su]
    [schema.core :as schema]
+   [toucan2.connection :as t2.conn]
    [toucan2.core :as t2]))
 
 (set! *warn-on-reflection* true)
@@ -45,17 +46,27 @@
 
 (declare premium-embedding-token)
 
-(defn- active-users-count* []
-  {:post [(integer? %)]}
-  "Returns the number of active users."
-  (assert ((requiring-resolve 'metabase.db/db-is-set-up?)) "Metabase DB is not yet set up")
-  (t2/count :core_user :is_active true))
-
-(def ^:private cached-active-user-count
-  "Primarily used for the settings because we don't wish it to be 100%."
-  (memoize/ttl
-    active-users-count*
-    :ttl/threshold (u/minutes->ms 5)))
+;; let's prevent the DB from getting slammed with calls to get the active user count, we only really need one in flight
+;; at a time.
+(let [f        (fn []
+                 {:post [(integer? %)]}
+                 (log/info (u/colorize :yellow "GETTING ACTIVE USER COUNT!"))
+                 (assert ((requiring-resolve 'metabase.db/db-is-set-up?)) "Metabase DB is not yet set up")
+                 ;; force this to use a new Connection, it seems to be getting called in situations where the Connection
+                 ;; is from a different thread and is invalid by the time we get to use it
+                 (let [result (binding [t2.conn/*current-connectable* nil]
+                                (t2/count :core_user :is_active true))]
+                   (log/info (u/colorize :green "=>") result)
+                   result))
+      memoized (memoize/ttl
+                f
+                :ttl/threshold (u/minutes->ms 5))
+      lock     (Object.)]
+  (defn- cached-active-users-count
+    "Primarily used for the settings because we don't wish it to be 100%. (HUH?)"
+    []
+    (locking lock
+      (memoized))))
 
 (defsetting active-users-count
   (deferred-tru "Cached number of active users. Refresh every 5 minutes.")
@@ -65,7 +76,7 @@
   :getter     (fn []
                 (if-not ((requiring-resolve 'metabase.db/db-is-set-up?))
                  0
-                 (cached-active-user-count))))
+                 (cached-active-users-count))))
 
 (defn- token-status-url [token base-url]
   (when (seq token)
@@ -86,7 +97,7 @@
 (defn- fetch-token-and-parse-body
   [token base-url]
   (some-> (token-status-url token base-url)
-          (http/get {:query-params {:users     (active-users-count*)
+          (http/get {:query-params {:users     (cached-active-users-count)
                                     :site-uuid (setting/get :site-uuid-for-premium-features-token-checks)}})
           :body
           (json/parse-string keyword)))
@@ -143,7 +154,7 @@
                 ;; tests to fail because a timed-out token check would get cached as a result.
                 (assert ((requiring-resolve 'metabase.db/db-is-set-up?)) "Metabase DB is not yet set up")
                 (u/with-timeout (u/seconds->ms 5)
-                  (active-users-count*))
+                  (cached-active-users-count))
                 (fetch-token-status* token))
               :ttl/threshold (u/minutes->ms 5))]
     (fn [token]

test/metabase/public_settings/premium_features_test.clj

@@ -40,7 +40,7 @@
   [token premium-features-response]
   (http-fake/with-fake-routes-in-isolation
     {{:address      (#'premium-features/token-status-url token @#'premium-features/token-check-url)
-      :query-params {:users     (str (#'premium-features/active-users-count*))
+      :query-params {:users     (str (#'premium-features/cached-active-users-count))
                      :site-uuid (public-settings/site-uuid-for-premium-features-token-checks)}}
      (constantly premium-features-response)}
     (#'premium-features/fetch-token-status* token)))
@@ -264,7 +264,8 @@
     ;; premium-features/active-users-count is cached so it could be make the test flaky
     ;; rebinding to avoid caching
     (testing "returns the number of active users"
-      (with-redefs [premium-features/cached-active-user-count #'premium-features/active-users-count*]
+      (with-redefs [premium-features/cached-active-users-count (fn []
+                                                                 (t2/count :core_user :is_active true))]
         (is (= (t2/count :core_user :is_active true)
                (premium-features/active-users-count)))))