Warm reboot: Add support for docker upgrade (sonic-net#292)

* Add support for docker upgrade
* Remove the extra spaces
* Make pending task check optional and add function support to get docker tag from docker meta label
* Use proc.returncode instead proc result for orchagent_restart_check
* Remove extra spaces

sonic_installer/main.py

@@ -173,6 +173,20 @@ def remove_image(image):
         run_command('grub-set-default --boot-directory=' + HOST_PATH + ' 0')
         click.echo('Image removed')
 
+# TODO: Embed tag name info into docker image meta data at build time,
+# and extract tag name from docker image file.
+def get_docker_tag_name(image):
+    # Try to get tag name from label metadata
+    cmd = "docker inspect --format '{{.ContainerConfig.Labels.Tag}}' " + image
+    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
+    (out, err) = proc.communicate()
+    if proc.returncode != 0:
+        return "unknown"
+    tag = out.rstrip()
+    if tag == "<no value>":
+        return "unknown"
+    return tag
+
 # Callback for confirmation prompt. Aborts if user enters "n"
 def abort_if_false(ctx, param, value):
     if not value:
@@ -343,5 +357,92 @@ def cleanup():
     if image_removed == 0:
         click.echo("No image(s) to remove")
 
+# Upgrade docker image
+@cli.command()
+@click.option('-y', '--yes', is_flag=True, callback=abort_if_false,
+        expose_value=False, prompt='New docker image will be installed, continue?')
+@click.option('--cleanup_image', is_flag=True, help="Clean up old docker image")
+@click.option('--enforce_check', is_flag=True, help="Enforce pending task check for docker upgrade")
+@click.option('--tag', type=str, help="Tag for the new docker image")
+@click.argument('container_name', metavar='<container_name>', required=True,
+    type=click.Choice(["swss", "snmp", "lldp", "bgp", "pmon", "dhcp_relay", "telemetry", "teamd"]))
+@click.argument('url')
+def upgrade_docker(container_name, url, cleanup_image, enforce_check, tag):
+    """ Upgrade docker image from local binary or URL"""
+
+    # example image: docker-lldp-sv2:latest
+    cmd = "docker inspect --format '{{.Config.Image}}' " + container_name
+    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
+    (out, err) = proc.communicate()
+    if proc.returncode != 0:
+        sys.exit(proc.returncode)
+    image_latest = out.rstrip()
+
+    # example image_name: docker-lldp-sv2
+    cmd = "echo " + image_latest + " | cut -d ':' -f 1"
+    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
+    image_name = proc.stdout.read().rstrip()
+
+    DEFAULT_IMAGE_PATH = os.path.join("/tmp/", image_name)
+    if url.startswith('http://') or url.startswith('https://'):
+        click.echo('Downloading image...')
+        try:
+            urllib.urlretrieve(url, DEFAULT_IMAGE_PATH, reporthook)
+        except Exception, e:
+            click.echo("Download error", e)
+            return
+        image_path = DEFAULT_IMAGE_PATH
+    else:
+        image_path = os.path.join("./", url)
+
+    # make sure orchagent is in clean state if swss is to be upgraded
+    if container_name == "swss":
+        skipPendingTaskCheck = " -s"
+        if enforce_check:
+            skipPendingTaskCheck = ""
+
+        cmd = "docker exec -it " + container_name + " orchagent_restart_check" + skipPendingTaskCheck
+        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
+        (out, err) = proc.communicate()
+        if proc.returncode != 0:
+            if enforce_check:
+                click.echo("Orchagent is not in clean state, check syslog and try later")
+                sys.exit(proc.returncode)
+            else:
+                click.echo("Orchagent is not in clean state, upgrading it anyway")
+        else:
+            click.echo("Orchagent is in clean state and frozen for warm upgrade")
+
+    run_command("systemctl stop %s" % container_name)
+    run_command("docker rm %s " % container_name)
+    run_command("docker rmi %s " % image_latest)
+    run_command("docker load < %s" % image_path)
+    if tag == None:
+        # example image: docker-lldp-sv2:latest
+        tag = get_docker_tag_name(image_latest)
+    run_command("docker tag %s:latest %s:%s" % (image_name, image_name, tag))
+    run_command("systemctl restart %s" % container_name)
+
+    # Clean up old docker images
+    if cleanup_image:
+        # All images id under the image name
+        cmd = "docker images --format '{{.ID}}' " + image_name
+        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
+        image_id_all = proc.stdout.read()
+        image_id_all = image_id_all.splitlines()
+        image_id_all = set(image_id_all)
+
+        # this is image_id for image with "latest" tag
+        cmd = "docker images --format '{{.ID}}' " + image_latest
+        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
+        image_id_latest = proc.stdout.read().rstrip()
+
+        for id in image_id_all:
+            if id != image_id_latest:
+                run_command("docker rmi -f %s" % id)
+
+    run_command("sleep 5") # wait 5 seconds for application to sync
+    click.echo('Done')
+
 if __name__ == '__main__':
     cli()