From b770e15a91d9f97216a06d766f5d3eba91cf5b9e Mon Sep 17 00:00:00 2001
From: "Joshua M. Boniface" <joshua@boniface.me>
Date: Sun, 18 Jul 2021 19:53:00 -0400
Subject: [PATCH] Fix final termination of logger

We need to do a bit more finagling with the logger on termination to
ensure that all messages are written and the queue drained before
actually terminating.
---
 daemon-common/log.py           | 36 +++++++++++++++++++++++++++-------
 node-daemon/pvcnoded/Daemon.py |  1 +
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/daemon-common/log.py b/daemon-common/log.py
index 6cb34f26..04e811b2 100644
--- a/daemon-common/log.py
+++ b/daemon-common/log.py
@@ -91,6 +91,17 @@ class Logger(object):
         self.writer.close()
         self.writer = open(self.logfile, 'a', buffering=0)
 
+    # Provide a termination function so all messages are flushed before terminating the main daemon
+    def terminate(self):
+        if self.config['file_logging']:
+            self.writer.close()
+        if self.config['zookeeper_logging']:
+            self.out("Waiting for Zookeeper message queue to drain", state='s')
+            while not self.zookeeper_logger.queue.empty():
+                pass
+            self.zookeeper_logger.stop()
+            self.zookeeper_logger.join()
+
     # Output function
     def out(self, message, state=None, prefix=''):
 
@@ -154,6 +165,7 @@ class ZookeeperLogger(Thread):
         self.zkhandler = None
         self.start_zkhandler()
         self.zkhandler.write([(('logs', self.node), '')])
+        self.running = False
         Thread.__init__(self, args=(), kwargs=None)
 
     def start_zkhandler(self):
@@ -167,18 +179,28 @@ class ZookeeperLogger(Thread):
         self.zkhandler.connect(persistent=True)
 
     def run(self):
-        while True:
-            message = self.queue.get()
-            self.write_log(message)
+        self.running = True
+        # Get the logs that are currently in Zookeeper and populate our deque
+        logs = deque(self.zkhandler.read(('logs.messages', self.node)).split('\n'), self.max_lines)
+        while self.running:
+            # Get a new message
+            try:
+                message = self.queue.get(block=False)
+                if not message:
+                    continue
+            except Exception:
+                continue
 
-    def write_log(self, message):
-        # Log to Zookeeper
-        with self.zkhandler.writelock(('logs.messages', self.node)):
-            logs = deque(self.zkhandler.read(('logs.messages', self.node)).split('\n'), self.max_lines)
             if not self.config['log_dates']:
                 # We want to log dates here, even if the log_dates config is not set
                 date = '{} '.format(datetime.now().strftime('%Y/%m/%d %H:%M:%S.%f'))
             else:
                 date = ''
+            # Add the message to the deque
             logs.append(f'{date}{message}')
+            # Write the updated messages into Zookeeper
             self.zkhandler.write([(('logs.messages', self.node), '\n'.join(logs))])
+        return
+
+    def stop(self):
+        self.running = False
diff --git a/node-daemon/pvcnoded/Daemon.py b/node-daemon/pvcnoded/Daemon.py
index 56fa6a65..1a13c947 100644
--- a/node-daemon/pvcnoded/Daemon.py
+++ b/node-daemon/pvcnoded/Daemon.py
@@ -761,6 +761,7 @@ def cleanup():
         pass
 
     logger.out('Terminated pvc daemon', state='s')
+    logger.terminate()
 
     os._exit(0)