Start edit monitor only when the feature is enabled

Check the env variable to determine whether the edit monitor feature is
enabled and only start the edit monitor if it is enabled, otherwise exit
directly. This helps to roll out the feature gradually.

Test: atest edit_monitor_utils_test atest daemon_manager_test
Bug: 365617369
Change-Id: I03fb494e1f62712efaf0bb05de8859e0118702bf

6 files changed, 224 insertions, 1 deletions

diff --git a/tools/edit_monitor/Android.bp b/tools/edit_monitor/Android.bp
index e613563153..b8ac5bff53 100644
--- a/tools/edit_monitor/Android.bp
+++ b/tools/edit_monitor/Android.bp
@@ -36,6 +36,7 @@ python_library_host {
     srcs: [
         "daemon_manager.py",
         "edit_monitor.py",
+        "utils.py",
     ],
     libs: [
         "asuite_cc_client",
@@ -75,6 +76,21 @@ python_test_host {
 }
 
 python_test_host {
+    name: "edit_monitor_utils_test",
+    main: "utils_test.py",
+    pkg_path: "edit_monitor",
+    srcs: [
+        "utils_test.py",
+    ],
+    libs: [
+        "edit_monitor_lib",
+    ],
+    test_options: {
+        unit_test: true,
+    },
+}
+
+python_test_host {
     name: "edit_monitor_integration_test",
     main: "edit_monitor_integration_test.py",
     pkg_path: "testdata",
diff --git a/tools/edit_monitor/daemon_manager.py b/tools/edit_monitor/daemon_manager.py
index c0a57abdaa..9a0abb6097 100644
--- a/tools/edit_monitor/daemon_manager.py
+++ b/tools/edit_monitor/daemon_manager.py
@@ -28,6 +28,7 @@ import time
 
 from atest.metrics import clearcut_client
 from atest.proto import clientanalytics_pb2
+from edit_monitor import utils
 from proto import edit_event_pb2
 
 DEFAULT_PROCESS_TERMINATION_TIMEOUT_SECONDS = 5
@@ -79,6 +80,15 @@ class DaemonManager:
 
   def start(self):
     """Writes the pidfile and starts the daemon proces."""
+    if not utils.is_feature_enabled(
+        "edit_monitor",
+        self.user_name,
+        "ENABLE_EDIT_MONITOR",
+        "EDIT_MONITOR_ROLLOUT_PERCENTAGE",
+    ):
+      logging.warning("Edit monitor is disabled, exiting...")
+      return
+
     if self.block_sign.exists():
       logging.warning("Block sign found, exiting...")
       return
diff --git a/tools/edit_monitor/daemon_manager_test.py b/tools/edit_monitor/daemon_manager_test.py
index e132000bce..407d94e9b5 100644
--- a/tools/edit_monitor/daemon_manager_test.py
+++ b/tools/edit_monitor/daemon_manager_test.py
@@ -81,6 +81,8 @@ class DaemonManagerTest(unittest.TestCase):
     # Sets the tempdir under the working dir so any temp files created during
     # tests will be cleaned.
     tempfile.tempdir = self.working_dir.name
+    self.patch = mock.patch.dict(os.environ, {'ENABLE_EDIT_MONITOR': 'true'})
+    self.patch.start()
 
   def tearDown(self):
     # Cleans up any child processes left by the tests.
@@ -88,6 +90,7 @@ class DaemonManagerTest(unittest.TestCase):
     self.working_dir.cleanup()
     # Restores tempdir.
     tempfile.tempdir = self.original_tempdir
+    self.patch.stop()
     super().tearDown()
 
   def test_start_success_with_no_existing_instance(self):
@@ -129,6 +132,15 @@ class DaemonManagerTest(unittest.TestCase):
 
     dm = daemon_manager.DaemonManager(TEST_BINARY_FILE)
     dm.start()
+
+    # Verify no daemon process is started.
+    self.assertIsNone(dm.daemon_process)
+
+  @mock.patch.dict(os.environ, {'ENABLE_EDIT_MONITOR': 'false'}, clear=True)
+  def test_start_return_directly_if_disabled(self):
+    dm = daemon_manager.DaemonManager(TEST_BINARY_FILE)
+    dm.start()
+
     # Verify no daemon process is started.
     self.assertIsNone(dm.daemon_process)
 
@@ -137,6 +149,7 @@ class DaemonManagerTest(unittest.TestCase):
         '/google/cog/cloud/user/workspace/edit_monitor'
     )
     dm.start()
+
     # Verify no daemon process is started.
     self.assertIsNone(dm.daemon_process)
 
diff --git a/tools/edit_monitor/edit_monitor_integration_test.py b/tools/edit_monitor/edit_monitor_integration_test.py
index d7dc7f1315..5f3d7e5a06 100644
--- a/tools/edit_monitor/edit_monitor_integration_test.py
+++ b/tools/edit_monitor/edit_monitor_integration_test.py
@@ -15,7 +15,6 @@
 """Integration tests for Edit Monitor."""
 
 import glob
-from importlib import resources
 import logging
 import os
 import pathlib
@@ -27,6 +26,9 @@ import tempfile
 import time
 import unittest
 
+from importlib import resources
+from unittest import mock
+
 
 class EditMonitorIntegrationTest(unittest.TestCase):
 
@@ -46,8 +48,11 @@ class EditMonitorIntegrationTest(unittest.TestCase):
     )
     self.root_monitoring_path.mkdir()
     self.edit_monitor_binary_path = self._import_executable("edit_monitor")
+    self.patch = mock.patch.dict(os.environ, {'ENABLE_EDIT_MONITOR': 'true'})
+    self.patch.start()
 
   def tearDown(self):
+    self.patch.stop()
     self.working_dir.cleanup()
     super().tearDown()
 
diff --git a/tools/edit_monitor/utils.py b/tools/edit_monitor/utils.py
new file mode 100644
index 0000000000..1a3275c6e2
--- /dev/null
+++ b/tools/edit_monitor/utils.py
@@ -0,0 +1,71 @@
+# Copyright 2024, The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import hashlib
+import logging
+import os
+
+
+def is_feature_enabled(
+    feature_name: str,
+    user_name: str,
+    enable_flag: str = None,
+    rollout_flag: str = None,
+) -> bool:
+  """Determine whether the given feature is enabled.
+
+  Whether a given feature is enabled or not depends on two flags: 1) the
+  enable_flag that explicitly enable/disable the feature and 2) the rollout_flag
+  that controls the rollout percentage.
+
+  Args:
+    feature_name: name of the feature.
+    user_name: system user name.
+    enable_flag: name of the env var that enables/disables the feature
+      explicitly.
+    rollout_flg: name of the env var that controls the rollout percentage, the
+      value stored in the env var should be an int between 0 and 100 string
+  """
+  if enable_flag:
+    if os.environ.get(enable_flag, "") == "false":
+      logging.info("feature: %s is disabled", feature_name)
+      return False
+
+    if os.environ.get(enable_flag, "") == "true":
+      logging.info("feature: %s is enabled", feature_name)
+      return True
+
+  if not rollout_flag:
+    return True
+
+  hash_object = hashlib.sha256()
+  hash_object.update((user_name + feature_name).encode("utf-8"))
+  hash_number = int(hash_object.hexdigest(), 16) % 100
+
+  roll_out_percentage = os.environ.get(rollout_flag, "0")
+  try:
+    percentage = int(roll_out_percentage)
+    if percentage < 0 or percentage > 100:
+      logging.warning(
+          "Rollout percentage: %s out of range, disable the feature.",
+          roll_out_percentage,
+      )
+      return False
+    return hash_number < percentage
+  except ValueError:
+    logging.warning(
+        "Invalid rollout percentage: %s, disable the feature.",
+        roll_out_percentage,
+    )
+    return False
diff --git a/tools/edit_monitor/utils_test.py b/tools/edit_monitor/utils_test.py
new file mode 100644
index 0000000000..7d7e4b207c
--- /dev/null
+++ b/tools/edit_monitor/utils_test.py
@@ -0,0 +1,108 @@
+# Copyright 2024, The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unittests for edit monitor utils."""
+import os
+import unittest
+from unittest import mock
+
+from edit_monitor import utils
+
+TEST_USER = 'test_user'
+TEST_FEATURE = 'test_feature'
+ENABLE_TEST_FEATURE_FLAG = 'ENABLE_TEST_FEATURE'
+ROLLOUT_TEST_FEATURE_FLAG = 'ROLLOUT_TEST_FEATURE'
+
+
+class EnableFeatureTest(unittest.TestCase):
+
+  def test_feature_enabled_without_flag(self):
+    self.assertTrue(utils.is_feature_enabled(TEST_FEATURE, TEST_USER))
+
+  @mock.patch.dict(os.environ, {ENABLE_TEST_FEATURE_FLAG: 'false'}, clear=True)
+  def test_feature_disabled_with_flag(self):
+    self.assertFalse(
+        utils.is_feature_enabled(
+            TEST_FEATURE, TEST_USER, ENABLE_TEST_FEATURE_FLAG
+        )
+    )
+
+  @mock.patch.dict(os.environ, {ENABLE_TEST_FEATURE_FLAG: 'true'}, clear=True)
+  def test_feature_enabled_with_flag(self):
+    self.assertTrue(
+        utils.is_feature_enabled(
+            TEST_FEATURE, TEST_USER, ENABLE_TEST_FEATURE_FLAG
+        )
+    )
+
+  @mock.patch.dict(
+      os.environ, {ROLLOUT_TEST_FEATURE_FLAG: 'invalid'}, clear=True
+  )
+  def test_feature_disabled_with_invalid_rollout_percentage(self):
+    self.assertFalse(
+        utils.is_feature_enabled(
+            TEST_FEATURE,
+            TEST_USER,
+            ENABLE_TEST_FEATURE_FLAG,
+            ROLLOUT_TEST_FEATURE_FLAG,
+        )
+    )
+
+  @mock.patch.dict(os.environ, {ROLLOUT_TEST_FEATURE_FLAG: '101'}, clear=True)
+  def test_feature_disabled_with_rollout_percentage_too_high(self):
+    self.assertFalse(
+        utils.is_feature_enabled(
+            TEST_FEATURE,
+            TEST_USER,
+            ENABLE_TEST_FEATURE_FLAG,
+            ROLLOUT_TEST_FEATURE_FLAG,
+        )
+    )
+
+  @mock.patch.dict(os.environ, {ROLLOUT_TEST_FEATURE_FLAG: '-1'}, clear=True)
+  def test_feature_disabled_with_rollout_percentage_too_low(self):
+    self.assertFalse(
+        utils.is_feature_enabled(
+            TEST_FEATURE,
+            TEST_USER,
+            ENABLE_TEST_FEATURE_FLAG,
+            ROLLOUT_TEST_FEATURE_FLAG,
+        )
+    )
+
+  @mock.patch.dict(os.environ, {ROLLOUT_TEST_FEATURE_FLAG: '90'}, clear=True)
+  def test_feature_enabled_with_rollout_percentage(self):
+    self.assertTrue(
+        utils.is_feature_enabled(
+            TEST_FEATURE,
+            TEST_USER,
+            ENABLE_TEST_FEATURE_FLAG,
+            ROLLOUT_TEST_FEATURE_FLAG,
+        )
+    )
+
+  @mock.patch.dict(os.environ, {ROLLOUT_TEST_FEATURE_FLAG: '10'}, clear=True)
+  def test_feature_disabled_with_rollout_percentage(self):
+    self.assertFalse(
+        utils.is_feature_enabled(
+            TEST_FEATURE,
+            TEST_USER,
+            ENABLE_TEST_FEATURE_FLAG,
+            ROLLOUT_TEST_FEATURE_FLAG,
+        )
+    )
+
+
+if __name__ == '__main__':
+  unittest.main()