AMBARI-11717. Ambari-agent died when trying to auto restart itself (dlysnichenko)

ambari-agent/src/main/python/ambari_agent/AlertSchedulerHandler.py

@@ -26,7 +26,6 @@ import logging
 import os
 import sys
 import time
-import atexit
 
 from apscheduler.scheduler import Scheduler
 from alerts.collector import AlertCollector
@@ -34,7 +33,7 @@ from alerts.metric_alert import MetricAlert
 from alerts.port_alert import PortAlert
 from alerts.script_alert import ScriptAlert
 from alerts.web_alert import WebAlert
-
+from ambari_agent.ExitHelper import ExitHelper
 logger = logging.getLogger(__name__)
 
 class AlertSchedulerHandler():
@@ -72,7 +71,7 @@ class AlertSchedulerHandler():
     self.config = config
 
     # register python exit handler
-    atexit.register(self.exit_handler)
+    ExitHelper().register(self.exit_handler)
 
 
   def exit_handler(self):

ambari-agent/src/main/python/ambari_agent/Controller.py

@@ -44,7 +44,7 @@ from ambari_agent.AlertSchedulerHandler import AlertSchedulerHandler
 from ambari_agent.ClusterConfiguration import  ClusterConfiguration
 from ambari_agent.RecoveryManager import  RecoveryManager
 from ambari_agent.HeartbeatHandlers import HeartbeatStopHandlers, bind_signal_handlers
-
+from ambari_agent.ExitHelper import ExitHelper
 logger = logging.getLogger(__name__)
 
 AGENT_AUTO_RESTART_EXIT_CODE = 77
@@ -384,7 +384,7 @@ class Controller(threading.Thread):
         self.heartbeatWithServer()
 
   def restartAgent(self):
-    sys.exit(AGENT_AUTO_RESTART_EXIT_CODE)
+    ExitHelper().exit(AGENT_AUTO_RESTART_EXIT_CODE)
 
 
   def sendRequest(self, url, data):

ambari-agent/src/main/python/ambari_agent/ExitHelper.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+
+'''
+Licensed to the Apache Software Foundation (ASF) under one
+or more contributor license agreements.  See the NOTICE file
+distributed with this work for additional information
+regarding copyright ownership.  The ASF licenses this file
+to you 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 os
+import logging
+import atexit
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["ExitHelper"]
+
+class _singleton(type):
+  _instances = {}
+
+  def __call__(cls, *args, **kwargs):
+    if cls not in cls._instances:
+      cls._instances[cls] = super(_singleton, cls).__call__(*args, **kwargs)
+    return cls._instances[cls]
+
+
+class ExitHelper(object):
+  """
+  Class to cleanup resources before exiting. Replacement for atexit module. sys.exit(code) works only from threads and
+  os._exit(code) will ignore atexit and cleanup will be ignored.
+  """
+  __metaclass__ = _singleton
+
+  def __init__(self):
+    self.exit_functions = []
+    self.exit_functions_executed = False
+    atexit.register(self.execute_cleanup)
+
+  def execute_cleanup(self):
+    if self.exit_functions_executed:
+      return
+    logger.info("Performing cleanup before exiting...")
+    while self.exit_functions:
+      func, args, kwargs = self.exit_functions.pop()
+      try:
+        func(*args, **kwargs)
+      except:
+        pass
+    self.exit_functions_executed = True
+
+  def register(self, func, *args, **kwargs):
+    self.exit_functions.append((func, args, kwargs))
+
+  def exit(self, code):
+    self.execute_cleanup()
+    os._exit(code)
+
+
+if __name__ == '__main__':
+  def func1():
+    print "1"
+
+  def func2():
+    print "2"
+
+  ExitHelper().register(func1)
+  ExitHelper().register(func2)
+  ExitHelper().exit(3)

ambari-agent/src/main/python/ambari_agent/apscheduler/threadpool.py

@@ -7,7 +7,7 @@ ThreadPool!
 from threading import Thread, Lock, currentThread
 from weakref import ref
 import logging
-import atexit
+from ambari_agent.ExitHelper import ExitHelper
 
 try:
     from queue import Queue, Empty
@@ -27,7 +27,7 @@ def _shutdown_all():
         if pool:
             pool.shutdown()
 
-atexit.register(_shutdown_all)
+ExitHelper().register(_shutdown_all)
 
 
 class ThreadPool(object):

ambari-agent/src/main/python/ambari_agent/main.py

@@ -36,6 +36,7 @@ from NetUtil import NetUtil
 from PingPortListener import PingPortListener
 import hostname
 from DataCleaner import DataCleaner
+from ExitHelper import ExitHelper
 import socket
 from ambari_commons import OSConst, OSCheck
 from ambari_commons.shell import shellRunner
@@ -278,6 +279,7 @@ def main(heartbeat_stop_callback=None):
     controller.start()
     controller.join()
   if not OSCheck.get_os_family() == OSConst.WINSRV_FAMILY:
+    ExitHelper.execute_cleanup()
     stop_agent()
   logger.info("finished")
 

ambari-agent/src/test/python/ambari_agent/TestController.py

@@ -35,6 +35,7 @@ from ambari_agent import hostname
 from ambari_agent.Controller import AGENT_AUTO_RESTART_EXIT_CODE
 from ambari_commons import OSCheck
 from ambari_agent.Hardware import Hardware
+from ambari_agent.ExitHelper import ExitHelper
 import ambari_commons
 
 OPERATING_SYSTEM_DISTRO = ('Suse','11','Final')
@@ -334,12 +335,12 @@ class TestController(unittest.TestCase):
     self.assertTrue(sendRequestMock.call_count > 5)
 
 
-  @patch("sys.exit")
-  def test_restartAgent(self, sys_exit_mock):
+  @patch.object(ExitHelper, "exit")
+  def test_restartAgent(self, exit_mock):
 
     self.controller.restartAgent()
-    self.assertTrue(sys_exit_mock.called)
-    self.assertTrue(sys_exit_mock.call_args[0][0] == AGENT_AUTO_RESTART_EXIT_CODE)
+    self.assertTrue(exit_mock.called)
+    self.assertTrue(exit_mock.call_args[0][0] == AGENT_AUTO_RESTART_EXIT_CODE)
 
 
   @patch("urllib2.Request")

ambari-agent/src/test/python/ambari_agent/TestMain.py

@@ -45,7 +45,7 @@ with patch.object(OSCheck, "os_distribution", new = MagicMock(return_value = os_
   from ambari_agent.DataCleaner import DataCleaner
   import ambari_agent.HeartbeatHandlers as HeartbeatHandlers
   from ambari_commons.os_check import OSConst, OSCheck
-
+  from ambari_agent.ExitHelper import ExitHelper
   if get_platform() != PLATFORM_WINDOWS:
     from ambari_commons.shell import shellRunnerLinux
 
@@ -306,7 +306,8 @@ class TestMain(unittest.TestCase):
   @patch.object(DataCleaner,"__init__")
   @patch.object(PingPortListener,"start")
   @patch.object(PingPortListener,"__init__")
-  def test_main(self, ping_port_init_mock, ping_port_start_mock, data_clean_init_mock,data_clean_start_mock,
+  @patch.object(ExitHelper,"execute_cleanup")
+  def test_main(self, cleanup_mock, ping_port_init_mock, ping_port_start_mock, data_clean_init_mock,data_clean_start_mock,
                 parse_args_mock, join_mock, start_mock, Controller_init_mock, try_to_connect_mock,
                 update_log_level_mock, daemonize_mock, perform_prestart_checks_mock,
                 ambari_config_mock,
@@ -339,7 +340,7 @@ class TestMain(unittest.TestCase):
     self.assertTrue(data_clean_start_mock.called)
     self.assertTrue(ping_port_init_mock.called)
     self.assertTrue(ping_port_start_mock.called)
-
+    self.assertTrue(cleanup_mock.called)
     perform_prestart_checks_mock.reset_mock()
 
     # Testing call with --expected-hostname parameter