From 9c914308f6f1376d5f11ae3875bb2be173d546c5 Mon Sep 17 00:00:00 2001
From: Darko Poljak <darko.poljak@gmail.com>
Date: Thu, 6 Jul 2017 12:43:17 +0200
Subject: [PATCH] Fix ssh connection multiplexing race condition #542

Increase ControlPersist to 2h.
After host run run ssh mux master exit command.
If custom remote exec/copy is specified then do nothing.
---
 cdist/__init__.py        |  2 +-
 cdist/config.py          | 37 +++++++++++++++++++++++++++++++++----
 cdist/util/remoteutil.py |  2 +-
 3 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/cdist/__init__.py b/cdist/__init__.py
index b0fd75ea..31e59794 100644
--- a/cdist/__init__.py
+++ b/cdist/__init__.py
@@ -20,7 +20,6 @@
 #
 
 import os
-import subprocess
 import hashlib
 
 import cdist.version
@@ -44,6 +43,7 @@ BANNER = """
 
 REMOTE_COPY = "scp -o User=root"
 REMOTE_EXEC = "ssh -o User=root"
+REMOTE_CMDS_CLEANUP_PATTERN = "ssh -o User=root -O exit -S {}"
 
 
 class Error(Exception):
diff --git a/cdist/config.py b/cdist/config.py
index 19679a4c..bd9d7bf8 100644
--- a/cdist/config.py
+++ b/cdist/config.py
@@ -47,13 +47,15 @@ from cdist.util.remoteutil import inspect_ssh_mux_opts
 class Config(object):
     """Cdist main class to hold arbitrary data"""
 
-    def __init__(self, local, remote, dry_run=False, jobs=None):
+    def __init__(self, local, remote, dry_run=False, jobs=None,
+                 cleanup_cmds=[]):
 
         self.local = local
         self.remote = remote
         self._open_logger()
         self.dry_run = dry_run
         self.jobs = jobs
+        self.cleanup_cmds = cleanup_cmds
 
         self.explorer = core.Explorer(self.local.target_host, self.local,
                                       self.remote, jobs=self.jobs)
@@ -91,6 +93,11 @@ class Config(object):
                 args.remote_exec_pattern = cdist.REMOTE_EXEC + mux_opts
             if args_dict['remote_copy'] is None:
                 args.remote_copy_pattern = cdist.REMOTE_COPY + mux_opts
+            if mux_opts:
+                cleanup_pattern = cdist.REMOTE_CMDS_CLEANUP_PATTERN
+            else:
+                cleanup_pattern = ""
+            args.remote_cmds_cleanup_pattern = cleanup_pattern
 
     @classmethod
     def _check_and_prepare_args(cls, args):
@@ -196,7 +203,12 @@ class Config(object):
             remote_copy = args.remote_copy_pattern.format(control_path)
         else:
             remote_copy = args.remote_copy
-        return (remote_exec, remote_copy, )
+        if args.remote_cmds_cleanup_pattern:
+            remote_cmds_cleanup = args.remote_cmds_cleanup_pattern.format(
+                control_path)
+        else:
+            remote_cmds_cleanup = ""
+        return (remote_exec, remote_copy, remote_cmds_cleanup, )
 
     @classmethod
     def onehost(cls, host, host_base_path, host_dir_name, args, parallel):
@@ -205,7 +217,8 @@ class Config(object):
         log = logging.getLogger(host)
 
         try:
-            remote_exec, remote_copy = cls._resolve_remote_cmds(args)
+            remote_exec, remote_copy, cleanup_cmd = cls._resolve_remote_cmds(
+                args)
             log.debug("remote_exec for host \"{}\": {}".format(
                 host, remote_exec))
             log.debug("remote_copy for host \"{}\": {}".format(
@@ -228,7 +241,11 @@ class Config(object):
                 remote_copy=remote_copy,
                 base_path=args.remote_out_path)
 
-            c = cls(local, remote, dry_run=args.dry_run, jobs=args.jobs)
+            cleanup_cmds = []
+            if cleanup_cmd:
+                cleanup_cmds.append(cleanup_cmd)
+            c = cls(local, remote, dry_run=args.dry_run, jobs=args.jobs,
+                    cleanup_cmds=cleanup_cmds)
             c.run()
 
         except cdist.Error as e:
@@ -272,11 +289,23 @@ class Config(object):
         self.explorer.run_global_explorers(self.local.global_explorer_out_path)
         self.manifest.run_initial_manifest(self.local.initial_manifest)
         self.iterate_until_finished()
+        self.cleanup()
 
         self.local.save_cache(start_time)
         self.log.info("Finished successful run in %s seconds",
                       time.time() - start_time)
 
+    def cleanup(self):
+        self.log.debug("Running cleanup commands")
+        for cleanup_cmd in self.cleanup_cmds:
+            cmd = cleanup_cmd.split()
+            cmd.append(self.local.target_host[0])
+            try:
+                self.local.run(cmd, return_output=False, save_output=False)
+            except cdist.Error as e:
+                # Log warning but continue.
+                self.log.warning("Cleanup command failed: %s", e)
+
     def object_list(self):
         """Short name for object list retrieval"""
         for cdist_object in core.CdistObject.list_objects(
diff --git a/cdist/util/remoteutil.py b/cdist/util/remoteutil.py
index 2bd12fdf..505c4598 100644
--- a/cdist/util/remoteutil.py
+++ b/cdist/util/remoteutil.py
@@ -36,7 +36,7 @@ def inspect_ssh_mux_opts():
     wanted_mux_opts = {
         "ControlPath": "{}",
         "ControlMaster": "auto",
-        "ControlPersist": "10",
+        "ControlPersist": "2h",
     }
     mux_opts = " ".join([" -o {}={}".format(
         x, wanted_mux_opts[x]) for x in wanted_mux_opts])