From 48db4bf8549638f9319cc346a8cdb0117493bb86 Mon Sep 17 00:00:00 2001 From: Mateusz Mandera Date: Thu, 9 Nov 2023 19:24:49 +0100 Subject: [PATCH] counts: Add new mobile_pushes RemoteRealmCount stats. This requires a bit of complexity to avoid a name collision in COUNT_STATS with the RemoteInstallationCount stats with the same name. --- analytics/lib/counts.py | 67 +++++-- .../commands/check_analytics_state.py | 4 +- .../management/commands/clear_single_stat.py | 4 +- .../commands/update_analytics_counts.py | 6 +- analytics/tests/test_counts.py | 174 +++++++++++++++++- zerver/lib/push_notifications.py | 1 + zerver/tests/test_push_notifications.py | 1 + ..._remoterealmcount_remote_realm_and_more.py | 25 +++ zilencer/models.py | 9 +- zilencer/views.py | 31 +++- 10 files changed, 293 insertions(+), 29 deletions(-) create mode 100644 zilencer/migrations/0035_remoterealmcount_remote_realm_and_more.py diff --git a/analytics/lib/counts.py b/analytics/lib/counts.py index 602c21b07c..cddc667887 100644 --- a/analytics/lib/counts.py +++ b/analytics/lib/counts.py @@ -24,7 +24,12 @@ from zerver.lib.timestamp import ceiling_to_day, ceiling_to_hour, floor_to_hour, from zerver.models import Message, Realm, RealmAuditLog, Stream, UserActivityInterval, UserProfile if settings.ZILENCER_ENABLED: - from zilencer.models import RemoteInstallationCount, RemoteZulipServer + from zilencer.models import ( + RemoteInstallationCount, + RemoteRealm, + RemoteRealmCount, + RemoteZulipServer, + ) ## Logging setup ## @@ -296,7 +301,7 @@ def do_aggregate_to_summary_table( # called from zerver.actions; should not throw any errors def do_increment_logging_stat( - model_object_for_bucket: Union[Realm, UserProfile, Stream, "RemoteZulipServer"], + model_object_for_bucket: Union[Realm, UserProfile, Stream, "RemoteRealm", "RemoteZulipServer"], stat: CountStat, subgroup: Optional[Union[str, int, bool]], event_time: datetime, @@ -308,9 +313,9 @@ def do_increment_logging_stat( table = stat.data_collector.output_table if table == RealmCount: assert isinstance(model_object_for_bucket, Realm) - id_args: Dict[str, Optional[Union[Realm, UserProfile, Stream, "RemoteZulipServer"]]] = { - "realm": model_object_for_bucket - } + id_args: Dict[ + str, Optional[Union[Realm, UserProfile, Stream, "RemoteRealm", "RemoteZulipServer"]] + ] = {"realm": model_object_for_bucket} elif table == UserCount: assert isinstance(model_object_for_bucket, UserProfile) id_args = {"realm": model_object_for_bucket.realm, "user": model_object_for_bucket} @@ -320,6 +325,13 @@ def do_increment_logging_stat( elif table == RemoteInstallationCount: assert isinstance(model_object_for_bucket, RemoteZulipServer) id_args = {"server": model_object_for_bucket, "remote_id": None} + elif table == RemoteRealmCount: + assert isinstance(model_object_for_bucket, RemoteRealm) + id_args = { + "server": model_object_for_bucket.server, + "remote_realm": model_object_for_bucket, + "remote_id": None, + } else: raise AssertionError("Unsupported CountStat output_table") @@ -850,25 +862,18 @@ def get_count_stats(realm: Optional[Realm] = None) -> Dict[str, CountStat]: ] if settings.ZILENCER_ENABLED: + # See also the remote_installation versions of these in REMOTE_INSTALLATION_COUNT_STATS. count_stats_.append( - # Tracks the number of push notifications requested to be sent - # by a remote server. LoggingCountStat( "mobile_pushes_received::day", - RemoteInstallationCount, + RemoteRealmCount, CountStat.DAY, ) ) count_stats_.append( - # Tracks the number of push notifications successfully sent to mobile - # devices, as requested by the remote server. Therefore this should be - # less than or equal to mobile_pushes_received - with potential tiny offsets - # resulting from a request being *received* by the bouncer right before midnight, - # but *sent* to the mobile device right after midnight. This would cause the increments - # to happen to CountStat records for different days. LoggingCountStat( "mobile_pushes_forwarded::day", - RemoteInstallationCount, + RemoteRealmCount, CountStat.DAY, ) ) @@ -886,3 +891,35 @@ BOUNCER_ONLY_REMOTE_COUNT_STAT_PROPERTIES = [ # To avoid refactoring for now COUNT_STATS can be used as before COUNT_STATS = get_count_stats() + +REMOTE_INSTALLATION_COUNT_STATS = OrderedDict() + +if settings.ZILENCER_ENABLED: + # REMOTE_INSTALLATION_COUNT_STATS contains duplicates of the + # RemoteRealmCount stats declared above; it is necessary because + # pre-8.0 servers do not send the fields required to identify a + # RemoteRealm. + + # Tracks the number of push notifications requested to be sent + # by a remote server. + REMOTE_INSTALLATION_COUNT_STATS["mobile_pushes_received::day"] = LoggingCountStat( + "mobile_pushes_received::day", + RemoteInstallationCount, + CountStat.DAY, + ) + # Tracks the number of push notifications successfully sent to + # mobile devices, as requested by the remote server. Therefore + # this should be less than or equal to mobile_pushes_received - + # with potential tiny offsets resulting from a request being + # *received* by the bouncer right before midnight, but *sent* to + # the mobile device right after midnight. This would cause the + # increments to happen to CountStat records for different days. + REMOTE_INSTALLATION_COUNT_STATS["mobile_pushes_forwarded::day"] = LoggingCountStat( + "mobile_pushes_forwarded::day", + RemoteInstallationCount, + CountStat.DAY, + ) + +ALL_COUNT_STATS = OrderedDict( + list(COUNT_STATS.items()) + list(REMOTE_INSTALLATION_COUNT_STATS.items()) +) diff --git a/analytics/management/commands/check_analytics_state.py b/analytics/management/commands/check_analytics_state.py index 606892ae82..6196dbace5 100644 --- a/analytics/management/commands/check_analytics_state.py +++ b/analytics/management/commands/check_analytics_state.py @@ -7,7 +7,7 @@ from django.core.management.base import BaseCommand from django.utils.timezone import now as timezone_now from typing_extensions import override -from analytics.lib.counts import COUNT_STATS, CountStat +from analytics.lib.counts import ALL_COUNT_STATS, CountStat from analytics.models import installation_epoch from zerver.lib.timestamp import TimeZoneNotUTCError, floor_to_day, floor_to_hour, verify_UTC from zerver.models import Realm @@ -44,7 +44,7 @@ class Command(BaseCommand): warning_unfilled_properties = [] critical_unfilled_properties = [] - for property, stat in COUNT_STATS.items(): + for property, stat in ALL_COUNT_STATS.items(): last_fill = stat.last_successful_fill() if last_fill is None: last_fill = installation_epoch() diff --git a/analytics/management/commands/clear_single_stat.py b/analytics/management/commands/clear_single_stat.py index c70a8ae66f..be7da9debc 100644 --- a/analytics/management/commands/clear_single_stat.py +++ b/analytics/management/commands/clear_single_stat.py @@ -4,7 +4,7 @@ from typing import Any from django.core.management.base import BaseCommand, CommandError from typing_extensions import override -from analytics.lib.counts import COUNT_STATS, do_drop_single_stat +from analytics.lib.counts import ALL_COUNT_STATS, do_drop_single_stat class Command(BaseCommand): @@ -18,7 +18,7 @@ class Command(BaseCommand): @override def handle(self, *args: Any, **options: Any) -> None: property = options["property"] - if property not in COUNT_STATS: + if property not in ALL_COUNT_STATS: raise CommandError(f"Invalid property: {property}") if not options["force"]: raise CommandError("No action taken. Use --force.") diff --git a/analytics/management/commands/update_analytics_counts.py b/analytics/management/commands/update_analytics_counts.py index 98aa81a0a4..20d136e4b6 100644 --- a/analytics/management/commands/update_analytics_counts.py +++ b/analytics/management/commands/update_analytics_counts.py @@ -10,7 +10,7 @@ from django.utils.dateparse import parse_datetime from django.utils.timezone import now as timezone_now from typing_extensions import override -from analytics.lib.counts import COUNT_STATS, logger, process_count_stat +from analytics.lib.counts import ALL_COUNT_STATS, logger, process_count_stat from scripts.lib.zulip_tools import ENDC, WARNING from zerver.lib.remote_server import send_analytics_to_push_bouncer from zerver.lib.timestamp import floor_to_hour @@ -74,9 +74,9 @@ class Command(BaseCommand): fill_to_time = floor_to_hour(fill_to_time.astimezone(timezone.utc)) if options["stat"] is not None: - stats = [COUNT_STATS[options["stat"]]] + stats = [ALL_COUNT_STATS[options["stat"]]] else: - stats = list(COUNT_STATS.values()) + stats = list(ALL_COUNT_STATS.values()) logger.info("Starting updating analytics counts through %s", fill_to_time) if options["verbose"]: diff --git a/analytics/tests/test_counts.py b/analytics/tests/test_counts.py index e1938cdd71..b1726978be 100644 --- a/analytics/tests/test_counts.py +++ b/analytics/tests/test_counts.py @@ -82,7 +82,13 @@ from zerver.models import ( get_user, is_cross_realm_bot_email, ) -from zilencer.models import RemoteInstallationCount, RemotePushDeviceToken, RemoteZulipServer +from zilencer.models import ( + RemoteInstallationCount, + RemotePushDeviceToken, + RemoteRealm, + RemoteRealmCount, + RemoteZulipServer, +) from zilencer.views import get_last_id_from_server @@ -244,7 +250,10 @@ class AnalyticsTestCase(ZulipTestCase): kwargs[arg_keys[i]] = values[i] for key, value in defaults.items(): kwargs[key] = kwargs.get(key, value) - if table not in [InstallationCount, RemoteInstallationCount] and "realm" not in kwargs: + if ( + table not in [InstallationCount, RemoteInstallationCount, RemoteRealmCount] + and "realm" not in kwargs + ): if "user" in kwargs: kwargs["realm"] = kwargs["user"].realm elif "stream" in kwargs: @@ -1433,6 +1442,9 @@ class TestLoggingCountStats(AnalyticsTestCase): hamlet, message, NotificationTriggers.DIRECT_MESSAGE ) + # First we'll make a request without providing realm_uuid. That means + # the bouncer can't increment the RemoteRealmCount stat, and only + # RemoteInstallationCount will be incremented. payload = { "user_id": hamlet.id, "user_uuid": str(hamlet.uuid), @@ -1466,8 +1478,162 @@ class TestLoggingCountStats(AnalyticsTestCase): RemoteInstallationCount, ["property", "value", "subgroup", "server", "remote_id", "end_time"], [ - ["mobile_pushes_received::day", 3, None, self.server, None, ceiling_to_day(now)], - ["mobile_pushes_forwarded::day", 2, None, self.server, None, ceiling_to_day(now)], + [ + "mobile_pushes_received::day", + 3, + None, + self.server, + None, + ceiling_to_day(now), + ], + [ + "mobile_pushes_forwarded::day", + 2, + None, + self.server, + None, + ceiling_to_day(now), + ], + ], + ) + self.assertFalse( + RemoteRealmCount.objects.filter(property="mobile_pushes_received::day").exists() + ) + self.assertFalse( + RemoteRealmCount.objects.filter(property="mobile_pushes_forwarded::day").exists() + ) + + # Now provide the realm_uuid. However, the RemoteRealm record doesn't exist yet, so it'll + # still be ignored. + payload = { + "user_id": hamlet.id, + "user_uuid": str(hamlet.uuid), + "realm_uuid": str(hamlet.realm.uuid), + "gcm_payload": gcm_payload, + "apns_payload": apns_payload, + "gcm_options": gcm_options, + } + with time_machine.travel(now, tick=False), mock.patch( + "zilencer.views.send_android_push_notification", return_value=1 + ), mock.patch( + "zilencer.views.send_apple_push_notification", return_value=1 + ), self.assertLogs( + "zilencer.views", level="INFO" + ): + result = self.uuid_post( + self.server_uuid, + "/api/v1/remotes/push/notify", + payload, + content_type="application/json", + subdomain="", + ) + self.assert_json_success(result) + + # The RemoteInstallationCount records get incremented again, but the RemoteRealmCount + # remains ignored due to missing RemoteRealm record. + self.assertTableState( + RemoteInstallationCount, + ["property", "value", "subgroup", "server", "remote_id", "end_time"], + [ + [ + "mobile_pushes_received::day", + 6, + None, + self.server, + None, + ceiling_to_day(now), + ], + [ + "mobile_pushes_forwarded::day", + 4, + None, + self.server, + None, + ceiling_to_day(now), + ], + ], + ) + self.assertFalse( + RemoteRealmCount.objects.filter(property="mobile_pushes_received::day").exists() + ) + self.assertFalse( + RemoteRealmCount.objects.filter(property="mobile_pushes_forwarded::day").exists() + ) + + # Create the RemoteRealm registration and repeat the above. This time RemoteRealmCount + # stats should be collected. + realm = hamlet.realm + remote_realm = RemoteRealm.objects.create( + server=self.server, + uuid=realm.uuid, + uuid_owner_secret=realm.uuid_owner_secret, + host=realm.host, + realm_deactivated=realm.deactivated, + realm_date_created=realm.date_created, + ) + + with time_machine.travel(now, tick=False), mock.patch( + "zilencer.views.send_android_push_notification", return_value=1 + ), mock.patch( + "zilencer.views.send_apple_push_notification", return_value=1 + ), self.assertLogs( + "zilencer.views", level="INFO" + ): + result = self.uuid_post( + self.server_uuid, + "/api/v1/remotes/push/notify", + payload, + content_type="application/json", + subdomain="", + ) + self.assert_json_success(result) + + # The RemoteInstallationCount records get incremented again, and the RemoteRealmCount + # gets collected. + self.assertTableState( + RemoteInstallationCount, + ["property", "value", "subgroup", "server", "remote_id", "end_time"], + [ + [ + "mobile_pushes_received::day", + 9, + None, + self.server, + None, + ceiling_to_day(now), + ], + [ + "mobile_pushes_forwarded::day", + 6, + None, + self.server, + None, + ceiling_to_day(now), + ], + ], + ) + self.assertTableState( + RemoteRealmCount, + ["property", "value", "subgroup", "server", "remote_realm", "remote_id", "end_time"], + [ + [ + "mobile_pushes_received::day", + 3, + None, + self.server, + remote_realm, + None, + ceiling_to_day(now), + ], + [ + "mobile_pushes_forwarded::day", + 2, + None, + self.server, + remote_realm, + None, + ceiling_to_day(now), + ], ], ) diff --git a/zerver/lib/push_notifications.py b/zerver/lib/push_notifications.py index aaa40020b0..51cb43f652 100644 --- a/zerver/lib/push_notifications.py +++ b/zerver/lib/push_notifications.py @@ -561,6 +561,7 @@ def send_notifications_to_bouncer( # user_uuid is the intended future format, but we also need to send user_id # to avoid breaking old mobile registrations, which were made with user_id. "user_id": user_profile.id, + "realm_uuid": str(user_profile.realm.uuid), "apns_payload": apns_payload, "gcm_payload": gcm_payload, "gcm_options": gcm_options, diff --git a/zerver/tests/test_push_notifications.py b/zerver/tests/test_push_notifications.py index fcc2bb5e56..731ce5a7be 100644 --- a/zerver/tests/test_push_notifications.py +++ b/zerver/tests/test_push_notifications.py @@ -2938,6 +2938,7 @@ class TestSendNotificationsToBouncer(ZulipTestCase): post_data = { "user_uuid": user.uuid, "user_id": user.id, + "realm_uuid": user.realm.uuid, "apns_payload": {"apns": True}, "gcm_payload": {"gcm": True}, "gcm_options": {}, diff --git a/zilencer/migrations/0035_remoterealmcount_remote_realm_and_more.py b/zilencer/migrations/0035_remoterealmcount_remote_realm_and_more.py new file mode 100644 index 0000000000..3a7b041b02 --- /dev/null +++ b/zilencer/migrations/0035_remoterealmcount_remote_realm_and_more.py @@ -0,0 +1,25 @@ +# Generated by Django 4.2.6 on 2023-11-09 17:23 + +import django.db.models.deletion +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("zilencer", "0034_remoterealmauditlog_remote_realm_and_more"), + ] + + operations = [ + migrations.AddField( + model_name="remoterealmcount", + name="remote_realm", + field=models.ForeignKey( + null=True, on_delete=django.db.models.deletion.CASCADE, to="zilencer.remoterealm" + ), + ), + migrations.AlterField( + model_name="remoterealmcount", + name="realm_id", + field=models.IntegerField(null=True), + ), + ] diff --git a/zilencer/models.py b/zilencer/models.py index 46440c1962..c90dba1d09 100644 --- a/zilencer/models.py +++ b/zilencer/models.py @@ -200,7 +200,14 @@ class RemoteInstallationCount(BaseRemoteCount): # We can't subclass RealmCount because we only have a realm_id here, not a foreign key. class RemoteRealmCount(BaseRemoteCount): - realm_id = models.IntegerField() + realm_id = models.IntegerField(null=True) + # Certain RemoteRealmCount will be counts tracked on the bouncer server directly, about + # stats pertaining to a realm on a remote server. For such objects, we will link to + # the corresponding RemoteRealm object that the remote server registered with us. + # In the future we may be able to link all RemoteRealmCount objects to a RemoteRealm, + # including the RemoteRealmCount objects are results of just syncing the RealmCount + # table from the remote server. + remote_realm = models.ForeignKey(RemoteRealm, on_delete=models.CASCADE, null=True) class Meta: unique_together = ("server", "realm_id", "property", "subgroup", "end_time") diff --git a/zilencer/views.py b/zilencer/views.py index 30e821aa03..c8fae7c8d3 100644 --- a/zilencer/views.py +++ b/zilencer/views.py @@ -20,6 +20,7 @@ from pydantic import BaseModel, ConfigDict from analytics.lib.counts import ( BOUNCER_ONLY_REMOTE_COUNT_STAT_PROPERTIES, COUNT_STATS, + REMOTE_INSTALLATION_COUNT_STATS, do_increment_logging_stat, ) from corporate.lib.stripe import do_deactivate_remote_server @@ -352,6 +353,16 @@ def remote_server_notify_push( apns_payload = payload["apns_payload"] gcm_options = payload.get("gcm_options", {}) + realm_uuid = payload.get("realm_uuid") + remote_realm = None + if realm_uuid is not None: + try: + remote_realm = RemoteRealm.objects.get(uuid=realm_uuid, server=server) + except RemoteRealm.DoesNotExist: + # We don't yet have a RemoteRealm for this realm. E.g. the server hasn't yet + # submitted analytics data since the realm's creation. + remote_realm = None + android_devices = list( RemotePushDeviceToken.objects.filter( user_identity.filter_q(), @@ -406,11 +417,19 @@ def remote_server_notify_push( ) do_increment_logging_stat( server, - COUNT_STATS["mobile_pushes_received::day"], + REMOTE_INSTALLATION_COUNT_STATS["mobile_pushes_received::day"], None, timezone_now(), increment=len(android_devices) + len(apple_devices), ) + if remote_realm is not None: + do_increment_logging_stat( + remote_realm, + COUNT_STATS["mobile_pushes_received::day"], + None, + timezone_now(), + increment=len(android_devices) + len(apple_devices), + ) # Truncate incoming pushes to 200, due to APNs maximum message # sizes; see handle_remove_push_notification for the version of @@ -446,11 +465,19 @@ def remote_server_notify_push( do_increment_logging_stat( server, - COUNT_STATS["mobile_pushes_forwarded::day"], + REMOTE_INSTALLATION_COUNT_STATS["mobile_pushes_forwarded::day"], None, timezone_now(), increment=android_successfully_delivered + apple_successfully_delivered, ) + if remote_realm is not None: + do_increment_logging_stat( + remote_realm, + COUNT_STATS["mobile_pushes_forwarded::day"], + None, + timezone_now(), + increment=android_successfully_delivered + apple_successfully_delivered, + ) return json_success( request,