Merge pull request #34889 from techhat/regpersist

Add load_reg and save_reg for Thorium

salt/returners/local_cache.py

@@ -23,6 +23,7 @@ import salt.utils.jid
 import salt.exceptions
 
 # Import 3rd-party libs
+import msgpack
 import salt.ext.six as six
 
 
@@ -478,3 +479,47 @@ def get_endtime(jid):
     with salt.utils.fopen(etpath, 'r') as etfile:
         endtime = etfile.read().strip('\n')
     return endtime
+
+
+def _reg_dir():
+    '''
+    Return the reg_dir for the given job id
+    '''
+    return os.path.join(__opts__['cachedir'], 'thorium')
+
+
+def save_reg(data):
+    '''
+    Save the register to msgpack files
+    '''
+    reg_dir = _reg_dir()
+    regfile = os.path.join(reg_dir, 'register')
+    try:
+        if not os.path.exists():
+            os.makedirs(reg_dir)
+    except OSError as exc:
+        if exc.errno == errno.EEXIST:
+            pass
+        else:
+            raise
+    try:
+        with salt.utils.fopen(regfile, 'a') as fh_:
+            msgpack.dump(data, fh_)
+        fh_.close()
+    except:
+        log.error('Could not write to msgpack file {0}'.format(__opts__['outdir']))
+        raise
+
+
+def load_reg():
+    '''
+    Load the register from msgpack files
+    '''
+    reg_dir = _reg_dir()
+    regfile = os.path.join(reg_dir, 'register')
+    try:
+        with salt.utils.fopen(regfile, 'r') as fh_:
+            return msgpack.load(fh_)
+    except:
+        log.error('Could not write to msgpack file {0}'.format(__opts__['outdir']))
+        raise

salt/thorium/__init__.py

@@ -18,6 +18,7 @@ import traceback
 
 # Import Salt libs
 import salt.state
+import salt.loader
 import salt.payload
 from salt.exceptions import SaltRenderError
 
@@ -43,7 +44,17 @@ class ThorState(salt.state.HighState):
         opts['file_client'] = 'local'
         self.opts = opts
         salt.state.HighState.__init__(self, self.opts, loader='thorium')
-        self.state.inject_globals = {'__reg__': {}}
+
+        self.returners = salt.loader.returners(self.opts, {})
+        self.reg_ret = self.opts.get('register_returner', None)
+        if self.reg_ret is not None:
+            try:
+                regdata = self.returners['{0}.load_reg'.format(self.reg_ret)]()
+            except Exception as exc:
+                log.error(exc)
+                regdata = {}
+
+        self.state.inject_globals = {'__reg__': regdata}
         self.event = salt.utils.event.get_master_event(
                 self.opts,
                 self.opts['sock_dir'])
@@ -174,4 +185,6 @@ class ThorState(salt.state.HighState):
             if (start - r_start) > recompile:
                 cache = self.gather_cache()
                 chunks = self.get_chunks()
+                if self.reg_ret is not None:
+                    self.returners['{0}.save_reg'.format(self.reg_ret)](chunks)
                 r_start = time.time()