[autotest] Consolidate methods required to setup a scheduler.
Move methods/classes that will be helpful in setting up another scheduler
process into scheduler_lib:
1. Make a connection manager capable of managing connections.
Create, access, close the database connection through this manager.
2. Cleanup setup_logging so it's usable by multiple schedulers if they
just change the name of the logfile.
TEST=Ran suites, unittests.
BUG=chromium:344613
DEPLOY=Scheduler
Change-Id: Id0031df96948d386416ce7cfc754f80456930b95
Reviewed-on: https://chromium-review.googlesource.com/199957
Reviewed-by: Prashanth B <beeps@chromium.org>
Tested-by: Prashanth B <beeps@chromium.org>
Commit-Queue: Prashanth B <beeps@chromium.org>
diff --git a/scheduler/host_scheduler.py b/scheduler/host_scheduler.py
index 38ff851..5ebdd7c 100644
--- a/scheduler/host_scheduler.py
+++ b/scheduler/host_scheduler.py
@@ -17,10 +17,6 @@
'get_metahost_schedulers', lambda : ())
-class SchedulerError(Exception):
- """Raised by HostScheduler when an inconsistent state occurs."""
-
-
class BaseHostScheduler(metahost_scheduler.HostSchedulingUtility):
"""Handles the logic for choosing when to run jobs and on which hosts.
diff --git a/scheduler/monitor_db.py b/scheduler/monitor_db.py
index ef305b4..31efe83 100755
--- a/scheduler/monitor_db.py
+++ b/scheduler/monitor_db.py
@@ -14,16 +14,16 @@
import django.db
-from autotest_lib.client.common_lib import global_config, logging_manager
+from autotest_lib.client.common_lib import global_config
from autotest_lib.client.common_lib import utils
-from autotest_lib.database import database_connection
-from autotest_lib.frontend.afe import models, rpc_utils, readonly_connection
+from autotest_lib.frontend.afe import models, rpc_utils
from autotest_lib.scheduler import agent_task, drone_manager, drones
from autotest_lib.scheduler import email_manager, gc_stats, host_scheduler
from autotest_lib.scheduler import monitor_db_cleanup, prejob_task
-from autotest_lib.scheduler import postjob_task, scheduler_logging_config
+from autotest_lib.scheduler import postjob_task
from autotest_lib.scheduler import rdb_lib
from autotest_lib.scheduler import rdb_utils
+from autotest_lib.scheduler import scheduler_lib
from autotest_lib.scheduler import scheduler_models
from autotest_lib.scheduler import status_server, scheduler_config
from autotest_lib.server import autoserv_utils
@@ -51,7 +51,7 @@
system error on the Autotest server. Full results may not be available. Sorry.
"""
-_db = None
+_db_manager = None
_shutdown = False
# These 2 globals are replaced for testing
@@ -75,7 +75,7 @@
def _verify_default_drone_set_exists():
if (models.DroneSet.drone_sets_enabled() and
not models.DroneSet.default_drone_set_name()):
- raise host_scheduler.SchedulerError(
+ raise scheduler_lib.SchedulerError(
'Drone sets are enabled, but no default is set')
@@ -98,8 +98,9 @@
def main_without_exception_handling():
- setup_logging()
-
+ scheduler_lib.setup_logging(
+ os.environ.get('AUTOTEST_SCHEDULER_LOG_DIR', None),
+ os.environ.get('AUTOTEST_SCHEDULER_LOG_NAME', None))
usage = 'usage: %prog [options] results_dir'
parser = optparse.OptionParser(usage)
parser.add_option('--recover-hosts', help='Try to recover dead hosts',
@@ -162,15 +163,7 @@
email_manager.manager.send_queued_emails()
server.shutdown()
_drone_manager.shutdown()
- _db.disconnect()
-
-
-def setup_logging():
- log_dir = os.environ.get('AUTOTEST_SCHEDULER_LOG_DIR', None)
- log_name = os.environ.get('AUTOTEST_SCHEDULER_LOG_NAME', None)
- logging_manager.configure_logging(
- scheduler_logging_config.SchedulerLoggingConfig(), log_dir=log_dir,
- logfile_name=log_name)
+ _db_manager.disconnect()
def handle_sigint(signum, frame):
@@ -193,14 +186,8 @@
DB_CONFIG_SECTION, 'database', 'stresstest_autotest_web')
os.environ['PATH'] = AUTOTEST_SERVER_DIR + ':' + os.environ['PATH']
- global _db
- _db = database_connection.DatabaseConnection(DB_CONFIG_SECTION)
- _db.connect(db_type='django')
-
- # ensure Django connection is in autocommit
- setup_django_environment.enable_autocommit()
- # bypass the readonly connection
- readonly_connection.ReadOnlyConnection.set_globally_disabled(True)
+ global _db_manager
+ _db_manager = scheduler_lib.ConnectionManager()
logging.info("Setting signal handler")
signal.signal(signal.SIGINT, handle_sigint)
@@ -248,11 +235,13 @@
def __init__(self):
self._agents = []
self._last_clean_time = time.time()
- self._host_scheduler = host_scheduler.HostScheduler(_db)
+ self._host_scheduler = host_scheduler.HostScheduler(
+ _db_manager.get_connection())
user_cleanup_time = scheduler_config.config.clean_interval
self._periodic_cleanup = monitor_db_cleanup.UserCleanup(
- _db, user_cleanup_time)
- self._24hr_upkeep = monitor_db_cleanup.TwentyFourHourUpkeep(_db)
+ _db_manager.get_connection(), user_cleanup_time)
+ self._24hr_upkeep = monitor_db_cleanup.TwentyFourHourUpkeep(
+ _db_manager.get_connection())
self._host_agents = {}
self._queue_entry_agents = {}
self._tick_count = 0
@@ -509,7 +498,7 @@
if queue_entry.status == models.HostQueueEntry.Status.ARCHIVING:
return postjob_task.ArchiveResultsTask(queue_entries=task_entries)
- raise host_scheduler.SchedulerError(
+ raise scheduler_lib.SchedulerError(
'_get_agent_task_for_queue_entry got entry with '
'invalid status %s: %s' % (queue_entry.status, queue_entry))
@@ -530,7 +519,7 @@
"""
if self.host_has_agent(entry.host):
agent = tuple(self._host_agents.get(entry.host.id))[0]
- raise host_scheduler.SchedulerError(
+ raise scheduler_lib.SchedulerError(
'While scheduling %s, host %s already has a host agent %s'
% (entry, entry.host, agent.task))
@@ -563,7 +552,7 @@
if agent_task_class.TASK_TYPE == special_task.task:
return agent_task_class(task=special_task)
- raise host_scheduler.SchedulerError(
+ raise scheduler_lib.SchedulerError(
'No AgentTask class for task', str(special_task))
@@ -629,7 +618,7 @@
if unrecovered_hqes:
message = '\n'.join(str(hqe) for hqe in unrecovered_hqes)
- raise host_scheduler.SchedulerError(
+ raise scheduler_lib.SchedulerError(
'%d unrecovered verifying host queue entries:\n%s' %
(len(unrecovered_hqes), message))
diff --git a/scheduler/monitor_db_unittest.py b/scheduler/monitor_db_unittest.py
index 5b12c91..62bf6b4 100755
--- a/scheduler/monitor_db_unittest.py
+++ b/scheduler/monitor_db_unittest.py
@@ -12,9 +12,10 @@
from autotest_lib.scheduler import agent_task
from autotest_lib.scheduler import monitor_db, drone_manager, email_manager
from autotest_lib.scheduler import pidfile_monitor
-from autotest_lib.scheduler import scheduler_config, gc_stats, host_scheduler
+from autotest_lib.scheduler import scheduler_config, gc_stats
from autotest_lib.scheduler import monitor_db_cleanup
from autotest_lib.scheduler import monitor_db_functional_test
+from autotest_lib.scheduler import scheduler_lib
from autotest_lib.scheduler import scheduler_models
_DEBUG = False
@@ -101,7 +102,9 @@
self._database.connect(db_type='django')
self._database.debug = _DEBUG
- self.god.stub_with(monitor_db, '_db', self._database)
+ connection_manager = scheduler_lib.ConnectionManager(autocommit=False)
+ self.god.stub_with(connection_manager, 'db_connection', self._database)
+ self.god.stub_with(monitor_db, '_db_manager', connection_manager)
self.god.stub_with(monitor_db.BaseDispatcher,
'_get_pending_queue_entries',
self._get_pending_hqes)
@@ -1021,7 +1024,7 @@
dummy_test_agent)
# Attempted to schedule on a host that already has an agent.
- self.assertRaises(host_scheduler.SchedulerError,
+ self.assertRaises(scheduler_lib.SchedulerError,
self._dispatcher._schedule_running_host_queue_entries)
diff --git a/scheduler/rdb_testing_utils.py b/scheduler/rdb_testing_utils.py
index 9600dfa..dab3a27 100644
--- a/scheduler/rdb_testing_utils.py
+++ b/scheduler/rdb_testing_utils.py
@@ -16,6 +16,7 @@
from autotest_lib.frontend.afe import rdb_model_extensions as rdb_models
from autotest_lib.scheduler import monitor_db
from autotest_lib.scheduler import monitor_db_functional_test
+from autotest_lib.scheduler import scheduler_lib
from autotest_lib.scheduler import scheduler_models
from autotest_lib.scheduler import rdb_hosts
from autotest_lib.scheduler import rdb_requests
@@ -236,7 +237,9 @@
self._database = self.db_helper.database
# Runs syncdb setting up initial database conditions
self._frontend_common_setup()
- self.god.stub_with(monitor_db, '_db', self._database)
+ connection_manager = scheduler_lib.ConnectionManager(autocommit=False)
+ self.god.stub_with(connection_manager, 'db_connection', self._database)
+ self.god.stub_with(monitor_db, '_db_manager', connection_manager)
self.god.stub_with(scheduler_models, '_db', self._database)
self._dispatcher = monitor_db.Dispatcher()
self.host_scheduler = self._dispatcher._host_scheduler
diff --git a/scheduler/scheduler_lib.py b/scheduler/scheduler_lib.py
new file mode 100644
index 0000000..5ccc9e7
--- /dev/null
+++ b/scheduler/scheduler_lib.py
@@ -0,0 +1,138 @@
+#pylint: disable-msg=C0111
+
+# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Scheduler helper libraries.
+"""
+import logging
+
+import common
+
+from autotest_lib.client.common_lib import logging_config
+from autotest_lib.client.common_lib import logging_manager
+from autotest_lib.database import database_connection
+from autotest_lib.frontend import setup_django_environment
+from autotest_lib.frontend.afe import readonly_connection
+
+
+DB_CONFIG_SECTION = 'AUTOTEST_WEB'
+
+
+class SchedulerError(Exception):
+ """Raised by the scheduler when an inconsistent state occurs."""
+
+
+class Singleton(type):
+ """Enforce that only one client class is instantiated per process."""
+ _instances = {}
+
+ def __call__(cls, *args, **kwargs):
+ """Fetch the instance of a class to use for subsequent calls."""
+ if cls not in cls._instances:
+ cls._instances[cls] = super(Singleton, cls).__call__(
+ *args, **kwargs)
+ return cls._instances[cls]
+
+
+class ConnectionManager(object):
+ """Manager for the django database connections.
+
+ The connection is used through scheduler_models and monitor_db.
+ """
+ __metaclass__ = Singleton
+
+ def __init__(self, readonly=False, autocommit=True):
+ """Set global django database options for correct connection handling.
+
+ @param readonly: Globally disable readonly connections.
+ @param autocommit: Initialize django autocommit options.
+ """
+ self.db_connection = None
+ # bypass the readonly connection
+ readonly_connection.ReadOnlyConnection.set_globally_disabled(readonly)
+ if autocommit:
+ # ensure Django connection is in autocommit
+ setup_django_environment.enable_autocommit()
+
+
+ @classmethod
+ def open_connection(cls):
+ """Open a new database connection.
+
+ @return: An instance of the newly opened connection.
+ """
+ db = database_connection.DatabaseConnection(DB_CONFIG_SECTION)
+ db.connect(db_type='django')
+ return db
+
+
+ def get_connection(self):
+ """Get a connection.
+
+ @return: A database connection.
+ """
+ if self.db_connection is None:
+ self.db_connection = self.open_connection()
+ return self.db_connection
+
+
+ def disconnect(self):
+ """Close the database connection."""
+ try:
+ self.db_connection.disconnect()
+ except Exception as e:
+ logging.debug('Could not close the db connection. %s', e)
+
+
+ def __del__(self):
+ self.disconnect()
+
+
+class SchedulerLoggingConfig(logging_config.LoggingConfig):
+ """Configure timestamped logging for a scheduler."""
+ GLOBAL_LEVEL = logging.INFO
+
+ @classmethod
+ def get_log_name(cls, timestamped_logfile_prefix):
+ """Get the name of a logfile.
+
+ @param timestamped_logfile_prefix: The prefix to apply to the
+ a timestamped log. Eg: 'scheduler' will create a logfile named
+ scheduler.log.2014-05-12-17.24.02.
+
+ @return: The timestamped log name.
+ """
+ return cls.get_timestamped_log_name(timestamped_logfile_prefix)
+
+
+ def configure_logging(self, log_dir=None, logfile_name=None,
+ timestamped_logfile_prefix='scheduler'):
+ """Configure logging to a specified logfile.
+
+ @param log_dir: The directory to log into.
+ @param logfile_name: The name of the log file.
+ @timestamped_logfile_prefix: The prefix to apply to the name of
+ the logfile, if a log file name isn't specified.
+ """
+ super(SchedulerLoggingConfig, self).configure_logging(use_console=True)
+
+ if log_dir is None:
+ log_dir = self.get_server_log_dir()
+ if not logfile_name:
+ logfile_name = self.get_log_name(timestamped_logfile_prefix)
+
+ self.add_file_handler(logfile_name, logging.DEBUG, log_dir=log_dir)
+
+
+def setup_logging(log_dir, log_name, timestamped_logfile_prefix='scheduler'):
+ """Setup logging to a given log directory and log file.
+
+ @param log_dir: The directory to log into.
+ @param log_name: Name of the log file.
+ @param timestamped_logfile_prefix: The prefix to apply to the logfile.
+ """
+ logging_manager.configure_logging(
+ SchedulerLoggingConfig(), log_dir=log_dir, logfile_name=log_name,
+ timestamped_logfile_prefix=timestamped_logfile_prefix)
diff --git a/scheduler/scheduler_lib_unittest.py b/scheduler/scheduler_lib_unittest.py
new file mode 100644
index 0000000..b8e3f84
--- /dev/null
+++ b/scheduler/scheduler_lib_unittest.py
@@ -0,0 +1,88 @@
+# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import mock
+import unittest
+
+import common
+
+from autotest_lib.database import database_connection
+from autotest_lib.frontend import setup_django_environment
+from autotest_lib.frontend.afe.readonly_connection import ReadOnlyConnection
+from autotest_lib.scheduler import scheduler_lib
+from django.db import utils as django_utils
+
+
+class ConnectionManagerTests(unittest.TestCase):
+ """Connection manager unittests."""
+
+ def setUp(self):
+ self.connection_manager = None
+ ReadOnlyConnection.set_globally_disabled = mock.MagicMock()
+ setup_django_environment.enable_autocommit = mock.MagicMock()
+ scheduler_lib.Singleton._instances = {}
+
+
+ def tearDown(self):
+ ReadOnlyConnection.set_globally_disabled.reset_mock()
+ setup_django_environment.enable_autocommit.reset_mock()
+
+
+ def testConnectionDisconnect(self):
+ """Test connection and disconnecting from the database."""
+ # Test that the connection manager only opens a connection once.
+ connection_manager = scheduler_lib.ConnectionManager()
+ connection_manager.open_connection = mock.MagicMock()
+ connection = connection_manager.get_connection()
+ connection_manager.open_connection.assert_called_once_with()
+ connection_manager.open_connection.reset_mock()
+ connection = connection_manager.get_connection()
+ self.assertTrue(
+ connection_manager.open_connection.call_count == 0)
+ connection_manager.open_connection.reset_mock()
+
+ # Test that del on the connection manager closes the connection
+ connection_manager.disconnect = mock.MagicMock()
+ connection_manager.__del__()
+ connection_manager.disconnect.assert_called_once_with()
+
+
+ def testConnectionReconnect(self):
+ """Test that retries don't destroy the connection."""
+ database_connection._DjangoBackend.execute = mock.MagicMock()
+ database_connection._DjangoBackend.execute.side_effect = (
+ django_utils.DatabaseError('Database Error'))
+ connection_manager = scheduler_lib.ConnectionManager()
+ connection = connection_manager.get_connection()
+ self.assertRaises(django_utils.DatabaseError,
+ connection.execute, *('', None, True))
+ self.assertTrue(
+ database_connection._DjangoBackend.execute.call_count == 2)
+ database_connection._DjangoBackend.execute.reset_mock()
+ self.assertTrue(connection_manager.db_connection ==
+ connection_manager.get_connection())
+
+
+ def testConnectionManagerSingleton(self):
+ """Test that the singleton works as expected."""
+ # Confirm that instantiating the class applies global db settings.
+ connection_manager = scheduler_lib.ConnectionManager()
+ ReadOnlyConnection.set_globally_disabled.assert_called_once_with(False)
+ setup_django_environment.enable_autocommit.assert_called_once_with()
+
+ ReadOnlyConnection.set_globally_disabled.reset_mock()
+ setup_django_environment.enable_autocommit.reset_mock()
+
+ # Confirm that instantiating another connection manager doesn't change
+ # the database settings, and in fact, returns the original manager.
+ connection_manager_2 = scheduler_lib.ConnectionManager()
+ self.assertTrue(connection_manager == connection_manager_2)
+ self.assertTrue(
+ ReadOnlyConnection.set_globally_disabled.call_count == 0)
+ self.assertTrue(
+ setup_django_environment.enable_autocommit.call_count == 0)
+
+ # Confirm that we don't open the connection when the class is
+ # instantiated.
+ self.assertTrue(connection_manager.db_connection is None)
diff --git a/scheduler/scheduler_logging_config.py b/scheduler/scheduler_logging_config.py
deleted file mode 100644
index eee9bbc..0000000
--- a/scheduler/scheduler_logging_config.py
+++ /dev/null
@@ -1,21 +0,0 @@
-import common
-import logging, os
-from autotest_lib.client.common_lib import logging_config
-
-class SchedulerLoggingConfig(logging_config.LoggingConfig):
- GLOBAL_LEVEL = logging.INFO
-
- @classmethod
- def get_log_name(cls):
- return cls.get_timestamped_log_name('scheduler')
-
-
- def configure_logging(self, log_dir=None, logfile_name=None):
- super(SchedulerLoggingConfig, self).configure_logging(use_console=True)
-
- if log_dir is None:
- log_dir = self.get_server_log_dir()
- if not logfile_name:
- logfile_name = self.get_log_name()
-
- self.add_file_handler(logfile_name, logging.DEBUG, log_dir=log_dir)
diff --git a/scheduler/scheduler_models.py b/scheduler/scheduler_models.py
index 1ebd7b8..2014108 100644
--- a/scheduler/scheduler_models.py
+++ b/scheduler/scheduler_models.py
@@ -22,10 +22,10 @@
from autotest_lib.client.common_lib import global_config, host_protections
from autotest_lib.client.common_lib import global_config, utils
from autotest_lib.frontend.afe import models, model_attributes
-from autotest_lib.database import database_connection
from autotest_lib.scheduler import drone_manager, email_manager
from autotest_lib.scheduler import rdb_lib
from autotest_lib.scheduler import scheduler_config
+from autotest_lib.scheduler import scheduler_lib
from autotest_lib.server.cros import provision
from autotest_lib.site_utils.graphite import stats
from autotest_lib.client.common_lib import control_data
@@ -38,8 +38,7 @@
def initialize():
global _db
- _db = database_connection.DatabaseConnection('AUTOTEST_WEB')
- _db.connect(db_type='django')
+ _db = scheduler_lib.ConnectionManager().get_connection()
notify_statuses_list = global_config.global_config.get_config_value(
scheduler_config.CONFIG_SECTION, "notify_email_statuses",
diff --git a/utils/unittest_suite.py b/utils/unittest_suite.py
index 8d46bd0..c9236d3 100755
--- a/utils/unittest_suite.py
+++ b/utils/unittest_suite.py
@@ -46,11 +46,13 @@
'rdb_unittest.py',
'rdb_hosts_unittest.py',
'rdb_cache_unittests.py',
+ 'scheduler_lib_unittest.py',
))
REQUIRES_MYSQLDB = set((
'migrate_unittest.py',
'db_utils_unittest.py',
+ 'scheduler_lib_unittest.py',
))
REQUIRES_GWT = set((
@@ -97,6 +99,7 @@
# crbug.com/251395
'dev_server_test.py',
'full_release_test.py',
+ 'scheduler_lib_unittest.py',
))
LONG_TESTS = (REQUIRES_MYSQLDB |