Island: Add Timeout and Version exceptions to mongo setup.

monkey/monkey_island/cc/server_setup.py

@@ -33,9 +33,11 @@ from monkey_island.cc.setup import island_config_options_validator  # noqa: E402
 from monkey_island.cc.setup.gevent_hub_error_handler import GeventHubErrorHandler  # noqa: E402
 from monkey_island.cc.setup.island_config_options import IslandConfigOptions  # noqa: E402
 from monkey_island.cc.setup.mongo.database_initializer import init_collections  # noqa: E402
-from monkey_island.cc.setup.mongo.mongo_db_process import MongoDbProcessException  # noqa: E402
 from monkey_island.cc.setup.mongo.mongo_setup import (  # noqa: E402
     MONGO_URL,
+    TIMEOUT,
+    MongoDBTimeOutException,
+    MongoDBVersionException,
     connect_to_mongodb,
     register_mongo_shutdown_callback,
     start_mongodb,
@@ -60,13 +62,17 @@ def run_monkey_island():
         register_mongo_shutdown_callback(mongo_db_process)
 
     try:
-        connect_to_mongodb(mongo_db_process)
-    except MongoDbProcessException:
-        logger.error(
-            f"MongoDB could not start. For details, check the MongoDB log at "
-            f"{mongo_db_process.get_log_file()}"
-        )
-        sys.exit(-1)
+        connect_to_mongodb()
+    except MongoDBTimeOutException:
+        if config_options.start_mongodb and not mongo_db_process.is_running():
+            logger.error(
+                f"Failed to start MongoDB process. Check log at {mongo_db_process.log_file}."
+            )
+        else:
+            logger.error(f"Failed to connect to MongoDB after {TIMEOUT} seconds. ")
+        sys.exit(1)
+    except MongoDBVersionException:
+        sys.exit(1)
 
     _configure_gevent_exception_handling(Path(config_options.data_dir))
     _start_island_server(island_args.setup_only, config_options)

monkey/monkey_island/cc/setup/mongo/mongo_db_process.py

@@ -29,9 +29,6 @@ class MongoDbProcess:
                 self._mongo_run_cmd, stderr=subprocess.STDOUT, stdout=log
             )
 
-        if not self.is_running():
-            raise MongoDbProcessException
-
         logger.info("MongoDB has been launched!")
 
     def stop(self):
@@ -57,9 +54,6 @@ class MongoDbProcess:
 
         return False
 
-    def get_log_file(self) -> str:
+    @property
+    def log_file(self) -> str:
         return self._log_file
-
-
-class MongoDbProcessException(Exception):
-    pass

monkey/monkey_island/cc/setup/mongo/mongo_setup.py

@@ -1,14 +1,13 @@
 import atexit
 import logging
 import os
-import sys
 import time
 
 from monkey_island.cc.database import get_db_version, is_db_server_up
 from monkey_island.cc.server_utils.file_utils import create_secure_directory
 from monkey_island.cc.setup.mongo import mongo_connector
 from monkey_island.cc.setup.mongo.mongo_connector import MONGO_DB_HOST, MONGO_DB_NAME, MONGO_DB_PORT
-from monkey_island.cc.setup.mongo.mongo_db_process import MongoDbProcess, MongoDbProcessException
+from monkey_island.cc.setup.mongo.mongo_db_process import MongoDbProcess
 
 DB_DIR_NAME = "db"
 MONGO_LOG_FILENAME = "mongodb.log"
@@ -17,6 +16,7 @@ MONGO_URL = os.environ.get(
     "mongodb://{0}:{1}/{2}".format(MONGO_DB_HOST, MONGO_DB_PORT, MONGO_DB_NAME),
 )
 MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.2.0"
+TIMEOUT = 15
 
 logger = logging.getLogger(__name__)
 
@@ -26,7 +26,6 @@ def start_mongodb(data_dir: str) -> MongoDbProcess:
     log_file = os.path.join(data_dir, MONGO_LOG_FILENAME)
 
     mongo_db_process = MongoDbProcess(db_dir=db_dir, log_file=log_file)
-
     mongo_db_process.start()
 
     return mongo_db_process
@@ -44,18 +43,20 @@ def register_mongo_shutdown_callback(mongo_db_process: MongoDbProcess):
     atexit.register(mongo_db_process.stop)
 
 
-def connect_to_mongodb(mongo_db_process: MongoDbProcess):
-    _wait_for_mongo_db_server(MONGO_URL, mongo_db_process)
+def connect_to_mongodb():
+    _wait_for_mongo_db_server(MONGO_URL)
     _assert_mongo_db_version(MONGO_URL)
     mongo_connector.connect_dal_to_mongodb()
 
 
-def _wait_for_mongo_db_server(mongo_url, mongo_db_process):
+def _wait_for_mongo_db_server(mongo_url):
+    start_time = time.time()
+
     while not is_db_server_up(mongo_url):
         logger.info("Waiting for MongoDB server on {0}".format(mongo_url))
 
-        if mongo_db_process is not None and not mongo_db_process.is_running():
-            raise MongoDbProcessException
+        if (time.time() - start_time) > TIMEOUT:
+            raise MongoDBTimeOutException
 
         time.sleep(1)
 
@@ -74,6 +75,14 @@ def _assert_mongo_db_version(mongo_url):
                 str(required_version), str(server_version)
             )
         )
-        sys.exit(-1)
+        raise MongoDBVersionException
     else:
         logger.info("Mongo DB version OK. Got {0}".format(str(server_version)))
+
+
+class MongoDBTimeOutException(Exception):
+    pass
+
+
+class MongoDBVersionException(Exception):
+    pass