[rearchitecture] Create function for doing MAB strategy selection. (#…

…3204)

Function was accidentally removed during a merge.
Related: #3008.

src/clusterfuzz/_internal/bot/tasks/utasks/fuzz_task.py

@@ -99,6 +99,26 @@ class FuzzErrorCode:
 Redzone = collections.namedtuple('Redzone', ['size', 'weight'])
 
 
+def do_multiarmed_bandit_strategy_selection(uworker_env):
+  """Set multi-armed bandit strategy selection during preprocessing. Set
+  multi-armed bandit strategy selection distribution as an environment variable
+  so we can access it in launcher."""
+  # TODO: Remove environment variable once fuzzing engine refactor is
+  # complete.
+  if not environment.get_value(
+      'USE_BANDIT_STRATEGY_SELECTION', env=uworker_env):
+    return
+  selection_method = utils.random_weighted_choice(SELECTION_METHOD_DISTRIBUTION,
+                                                  'probability')
+  environment.set_value('STRATEGY_SELECTION_METHOD',
+                        selection_method.method_name, uworker_env)
+  distribution = get_strategy_distribution_from_ndb()
+  if not distribution:
+    return
+  environment.set_value('STRATEGY_SELECTION_DISTRIBUTION', distribution,
+                        uworker_env)
+
+
 def get_unsymbolized_crash_stacktrace(stack_file_path):
   """Read unsymbolized crash stacktrace."""
   with open(stack_file_path, 'rb') as f:

src/clusterfuzz/_internal/system/environment.py

@@ -601,9 +601,11 @@ def get_value_string(environment_variable, default_value=None):
   return os.getenv(environment_variable, default_value)
 
 
-def get_value(environment_variable, default_value=None):
+def get_value(environment_variable, default_value=None, env=None):
   """Return an environment variable value."""
-  value_string = os.getenv(environment_variable)
+  if env is None:
+    env = os.environ
+  value_string = env.get(environment_variable)
 
   # value_string will be None if the variable is not defined.
   if value_string is None: