more WIP to special-case trial users

1 files changed, 55 insertions, 20 deletions

diff --git a/synapse/storage/monthly_active_users.py b/synapse/storage/monthly_active_users.py
index 03526ce1e3..6755fec389 100644
--- a/synapse/storage/monthly_active_users.py
+++ b/synapse/storage/monthly_active_users.py
@@ -46,7 +46,7 @@ class MonthlyActiveUsersStore(SQLBaseStore):
                 tp["medium"], tp["address"]
             )
             if user_id:
-                yield self.upsert_monthly_active_user(user_id)
+                yield self.upsert_monthly_active_user(user_id, False)
                 reserved_user_list.append(user_id)
             else:
                 logger.warning(
@@ -88,14 +88,36 @@ class MonthlyActiveUsersStore(SQLBaseStore):
 
             txn.execute(sql, query_args)
 
-            # If MAU user count still exceeds the MAU threshold, then delete on
-            # a least recently active basis.
+            # Promote trial users to non-trial users, oldest first, assuming we
+            # have MAU headroom available.  Otherwise we leave them stuck in trial
+            # purgatory until their 30 days is up.
+            #
+            # We don't need to worry about reserved users, as they are already non-trial.
+            mau_trial_ms = self.hs.config.mau_trial_days * 24 * 60 * 60 * 1000
+
+            sql = """
+                UPDATE monthly_active_users SET trial='n' WHERE user_id IN (
+                    SELECT user_id FROM monthly_active_users
+                    ORDER BY (timestamp - last_active) DESC
+                    WHERE trial='y'
+                    LIMIT ? - (SELECT count(*) FROM monthly_active_users WHERE trial='n')
+                ) AND timestamp - last_active >= ?
+                """
+
+            # FIXME: handle negative limits
+
+            txn.execute(sql, (self.hs.config.max_mau_value, mau_trial_ms))
+
+            # If non-trial MAU user count still exceeds the MAU threshold, then
+            # delete on a least recently active basis.
+            #
             # Note it is not possible to write this query using OFFSET due to
             # incompatibilities in how sqlite and postgres support the feature.
             # sqlite requires 'LIMIT -1 OFFSET ?', the LIMIT must be present
             # While Postgres does not require 'LIMIT', but also does not support
             # negative LIMIT values. So there is no way to write it that both can
             # support
+
             safe_guard = self.hs.config.max_mau_value - len(self.reserved_users)
             # Must be greater than zero for postgres
             safe_guard = safe_guard if safe_guard > 0 else 0
@@ -106,8 +128,9 @@ class MonthlyActiveUsersStore(SQLBaseStore):
                 WHERE user_id NOT IN (
                     SELECT user_id FROM monthly_active_users
                     ORDER BY timestamp DESC
+                    WHERE trial='n'
                     LIMIT ?
-                    )
+                    ) AND trial='n'
                 """
             # Need if/else since 'AND user_id NOT IN ({})' fails on Postgres
             # when len(reserved_users) == 0. Works fine on sqlite.
@@ -139,25 +162,24 @@ class MonthlyActiveUsersStore(SQLBaseStore):
             Defered[int]: Number of current monthly active users
         """
 
-        mau_trial_ms = self.hs.config.mau_trial_days * 24 * 60 * 60 * 1000
-
         def _count_users(txn):
             sql = """
                 SELECT COALESCE(count(*), 0)
                 FROM monthly_active_users
-                WHERE timestamp - first_active >= ?
+                WHERE trial = 'n'
             """
 
-            txn.execute(sql, (mau_trial_ms,))
+            txn.execute(sql)
             count, = txn.fetchone()
             return count
         return self.runInteraction("count_users", _count_users)
 
-    def upsert_monthly_active_user(self, user_id):
+    def upsert_monthly_active_user(self, user_id, trial):
         """
             Updates or inserts monthly active user member
             Arguments:
                 user_id (str): user to add/update
+                trial (bool): whether the user is entering a trial or not
             Deferred[bool]: True if a new entry was created, False if an
                 existing one was updated.
         """
@@ -173,6 +195,7 @@ class MonthlyActiveUsersStore(SQLBaseStore):
             },
             insertion_values={
                 "first_active": now,
+                "trial": trial,
             },
             lock=False,
         )
@@ -184,6 +207,7 @@ class MonthlyActiveUsersStore(SQLBaseStore):
     def user_last_seen_monthly_active(self, user_id):
         """
             Checks if a given user is part of the monthly active user group
+            or a trial user.
             Arguments:
                 user_id (str): user to add/update
             Return:
@@ -191,15 +215,26 @@ class MonthlyActiveUsersStore(SQLBaseStore):
 
         """
 
-        return(self._simple_select_one_onecol(
-            table="monthly_active_users",
-            keyvalues={
-                "user_id": user_id,
-            },
-            retcol="timestamp",
-            allow_none=True,
-            desc="user_last_seen_monthly_active",
-        ))
+        # FIXME: we should probably return whether this is a trial user or not.
+
+        mau_trial_ms = self.hs.config.mau_trial_days * 24 * 60 * 60 * 1000
+
+        def _user_last_seen_monthly_active_txn(txn):
+            sql = """
+                SELECT timestamp
+                FROM monthly_active_users
+                WHERE trial = 'n' OR (
+                    timestamp - last_active < ?
+                )
+            """
+
+            txn.execute(sql, (mau_trial_ms, ))
+            count, = txn.fetchone()
+            return count
+        return self.runInteraction(
+            "user_last_seen_monthly_active",
+            _user_last_seen_monthly_active
+        )
 
     @defer.inlineCallbacks
     def populate_monthly_active_users(self, user_id):
@@ -221,6 +256,6 @@ class MonthlyActiveUsersStore(SQLBaseStore):
             if last_seen_timestamp is None:
                 count = yield self.get_monthly_active_count()
                 if count < self.hs.config.max_mau_value:
-                    yield self.upsert_monthly_active_user(user_id)
+                    yield self.upsert_monthly_active_user(user_id, True)
             elif now - last_seen_timestamp > LAST_SEEN_GRANULARITY:
-                yield self.upsert_monthly_active_user(user_id)
+                yield self.upsert_monthly_active_user(user_id, True)