From 2659a0b3b945bcbf82605a3ac0384919ee73676b Mon Sep 17 00:00:00 2001
From: Hemna <waboring@hemna.com>
Date: Wed, 30 Dec 2020 07:32:20 -0500
Subject: [PATCH] Added support to save/load MsgTrack on exit/start

This patch added saving of the MsgTrack list of messages at aprsd exit.
The will be loaded at startup unless you pass in the --flush option.
---
 aprsd/main.py      | 28 ++++++++++++++++++++++++----
 aprsd/messaging.py | 39 ++++++++++++++++++++++++++++++++++-----
 aprsd/utils.py     |  6 +++++-
 3 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/aprsd/main.py b/aprsd/main.py
index a741477..416c996 100644
--- a/aprsd/main.py
+++ b/aprsd/main.py
@@ -150,7 +150,6 @@ def signal_handler(signal, frame):
     )
     threads.APRSDThreadList().stop_all()
     server_event.set()
-    sys.exit(0)  # thread ignores this
 
 
 # end signal_handler
@@ -333,7 +332,16 @@ def send_message(
     default=utils.DEFAULT_CONFIG_FILE,
     help="The aprsd config file to use for options.",
 )
-def server(loglevel, quiet, disable_validation, config_file):
+@click.option(
+    "-f",
+    "--flush",
+    "flush",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Flush out all old aged messages on disk.",
+)
+def server(loglevel, quiet, disable_validation, config_file, flush):
     """Start the aprsd server process."""
     global event
 
@@ -364,6 +372,15 @@ def server(loglevel, quiet, disable_validation, config_file):
     plugin_manager.setup_plugins()
     client.Client(config)
 
+    # Now load the msgTrack from disk if any
+    if flush:
+        LOG.debug("Deleting saved MsgTrack.")
+        messaging.MsgTrack().flush()
+    else:
+        # Try and load saved MsgTrack list
+        LOG.debug("Loading saved MsgTrack object.")
+        messaging.MsgTrack().load()
+
     rx_msg_queue = queue.Queue(maxsize=20)
     tx_msg_queue = queue.Queue(maxsize=20)
     msg_queues = {
@@ -378,6 +395,8 @@ def server(loglevel, quiet, disable_validation, config_file):
     rx_thread.start()
     tx_thread.start()
 
+    messaging.MsgTrack().restart()
+
     cntr = 0
     while not server_event.is_set():
         # to keep the log noise down
@@ -387,9 +406,10 @@ def server(loglevel, quiet, disable_validation, config_file):
         cntr += 1
         time.sleep(10)
 
+    # If there are items in the msgTracker, then save them
+    tracker = messaging.MsgTrack()
+    tracker.save()
     LOG.info("APRSD Exiting.")
-    sys.exit(0)
-    # setup and run the main blocking loop
 
 
 if __name__ == "__main__":
diff --git a/aprsd/messaging.py b/aprsd/messaging.py
index faed185..043e599 100644
--- a/aprsd/messaging.py
+++ b/aprsd/messaging.py
@@ -1,12 +1,15 @@
 import abc
 import datetime
 import logging
+import os
+import pathlib
+import pickle
 import re
 import threading
 import time
 from multiprocessing import RawValue
 
-from aprsd import client, threads
+from aprsd import client, threads, utils
 
 LOG = logging.getLogger("APRSD")
 
@@ -44,6 +47,7 @@ class MsgTrack(object):
     lock = None
 
     track = {}
+    total_messages_tracked = 0
 
     def __new__(cls, *args, **kwargs):
         if cls._instance is None:
@@ -56,6 +60,7 @@ class MsgTrack(object):
         with self.lock:
             key = int(msg.id)
             self.track[key] = msg
+            self.total_messages_tracked += 1
 
     def get(self, id):
         with self.lock:
@@ -82,11 +87,35 @@ class MsgTrack(object):
 
     def save(self):
         """Save this shit to disk?"""
-        pass
+        if len(self) > 0:
+            LOG.info("Need to save tracking to disk")
+            pickle.dump(self.dump(), open(utils.DEFAULT_SAVE_FILE, "wb+"))
 
-    def restore(self):
-        """Restore this shit?"""
-        pass
+    def dump(self):
+        dump = {}
+        with self.lock:
+            for key in self.track.keys():
+                dump[key] = self.track[key]
+
+        return dump
+
+    def load(self):
+        if os.path.exists(utils.DEFAULT_SAVE_FILE):
+            raw = pickle.load(open(utils.DEFAULT_SAVE_FILE, "rb"))
+            if raw:
+                self.track = raw
+                LOG.debug("Loaded MsgTrack dict from disk.")
+                LOG.debug(self)
+
+    def restart(self):
+        """Walk the list of messages and restart them if any."""
+        for key in self.track.keys():
+            msg = self.track[key]
+            msg.send()
+
+    def flush(self):
+        """Nuke the old pickle file that stored the old results from last aprsd run."""
+        pathlib.Path(utils.DEFAULT_SAVE_FILE).unlink()
 
 
 class MessageCounter(object):
diff --git a/aprsd/utils.py b/aprsd/utils.py
index 6d7970b..f92c7d9 100644
--- a/aprsd/utils.py
+++ b/aprsd/utils.py
@@ -5,6 +5,7 @@ import functools
 import os
 import sys
 import threading
+from pathlib import Path
 
 import click
 import yaml
@@ -46,7 +47,10 @@ DEFAULT_CONFIG_DICT = {
     },
 }
 
-DEFAULT_CONFIG_FILE = "~/.config/aprsd/aprsd.yml"
+home = str(Path.home())
+DEFAULT_CONFIG_DIR = "{}/.config/aprsd/".format(home)
+DEFAULT_SAVE_FILE = "{}/.config/aprsd/aprsd.p".format(home)
+DEFAULT_CONFIG_FILE = "{}/.config/aprsd/aprsd.yml".format(home)
 
 
 def synchronized(wrapped):