#!/bin/sh
# Example of how to use opengl with lxc and ubuntu 16.04, 18.04, or 19.04
#
# lxc 3.6 and up make this really easy (modulo a few odd nvidia settings).
#
# It's much simpler than this example makes it look, but
# when you're just starting out, sometimes
# having a full working example is more helpful than
# simple fragments.

set -e

# Which version this machine is running
case $(cat /etc/issue) in
*16.04*) _os=16;;
*18.04*) _os=18;;
*19.04*) _os=19;;
*) echo "Sorry, this demo has only been tested on Ubuntu 16.04, 18.04, and 19.04."; exit 1;;
esac

PATH=/snap/bin:$PATH

# Work around unreliable 'lxc stop', which hangs frequently
# https://discuss.linuxcontainers.org/t/lxc-stop-hangs-and-works-after-second-try/3598
lxc_stop() {
  lxc exec $1 -- shutdown -h now
}

usage() {
    cat << _EOF_
Demonstrate opengl on multiple flavors of Ubuntu in lxc.

Wipes your lxc/lxd configuration, so only run this on machines that do
not currently use lxc/lxd.

You'll need to do 'xhost +' locally on all machines first, since this
demo hasn't figured out X authentication yet.

Multimachine example:
to prepare ubuntu 16.04/18.04/19.04 images locally, then
set up lxc / run glxgears in lxc / tear down lxc
on two remote machines (running any of those three ubuntus):
  sh lxc-ubuntu-opengl-demo.sh serve 16 18 19
  sh lxc-ubuntu-opengl-demo.sh multi-setup host1 host2
  sh lxc-ubuntu-opengl-demo.sh multi-run host1 host2
  sh lxc-ubuntu-opengl-demo.sh multi-teardown host1 host2

Single-machine example: to prepare image and run glxinfo locally, run these steps, in order:
  sh lxc-ubuntu-opengl-demo.sh install
  sh lxc-ubuntu-opengl-demo.sh init
  sh lxc-ubuntu-opengl-demo.sh image
  sh lxc-ubuntu-opengl-demo.sh localize
  sh lxc-ubuntu-opengl-demo.sh run
  sh lxc-ubuntu-opengl-demo.sh purge
  sh lxc-ubuntu-opengl-demo.sh uninstall
_EOF_
}

# Uninstall lxd
do_uninstall() {
  if test -d /snap/lxd
  then
    sudo snap remove lxd
  fi
}

# Purge all traces of lxd's data from the universe
do_purge() {
  rm -f empty-lxd.sh
  wget https://raw.githubusercontent.com/lxc/lxd/master/scripts/empty-lxd.sh
  if ! test -x /usr/bin/jq
  then
    sudo apt install -y jq
  fi
  sh empty-lxd.sh
}

# Install the right version of lxd
# Idempotent.  Runs quickly and does nothing second time.
do_install() {
  # Can't use lxc from apt, it's waaay too old on some systems (ubuntu 16.04)
  if test -x /usr/bin/lxc
  then
    sudo apt remove --purge lxd lxd-client || true
  fi

  if ! snap --version
  then
    # what, you don't have snap installed already?
    sudo apt install snapd
  fi

  # Install lxc 3.0 from snap
  if ! test -d /snap/lxd
  then
    sudo snap install lxd
    echo "Note: you may need to log out and log back in now to get access to lxd from commandline."
  fi
  which lxc
  lxc --version
}

do_uninstall() {
  sudo snap remove lxd
}

# Initialize lxd, if not already initialized
# Idempotent.  Runs quickly and does nothing second time.
# You may need to add yourself to the lxd group before doing this.
do_init() {
  # fixme: handle more types of storage?
  if df -T /var/lib/lxd | grep btrfs
  then
    driver=btrfs
  else
    driver=dir
  fi

  if lxc storage ls | grep default
  then
    echo "lxd is already initialized"
  else
    lxd init --preseed <<_EOF_
config: {}
networks:
- config:
    ipv4.address: auto
    ipv6.address: none
  description: ""
  managed: false
  name: lxdbr0
  type: ""
storage_pools:
- config:
    source: /var/lib/lxd/storage-pools/default
  description: ""
  name: default
  driver: $driver
profiles:
- config: {}
  description: ""
  devices:
    eth0:
      name: eth0
      nictype: bridged
      parent: lxdbr0
      type: nic
    root:
      path: /
      pool: default
      type: disk
  name: default
cluster: null
_EOF_
  fi
}

# Create Ubuntu images ready to run an opengl application
# and publish them with alias gpu-${OS}
# Usage: do_image [releaseyear...]
# Default is same release as on local system.
do_image() {
  OSES="$*"
  # Default to local OS
  case "$OSES" in
  "") OSES=$_os;; 
  esac

  for OS in $OSES
  do
    # Prepare a base image gpu-${OS} with updates and a few packages, for fast startup
    lxc_stop preloaded-${OS} || true
    # fixme: make deletion optional; otherwise just return if image already there
    lxc delete preloaded-${OS} || true
    lxc launch --profile default ubuntu:${OS}.04 preloaded-${OS}
    # Wait for network to come up
    while ! lxc exec preloaded-${OS} -- host archive.ubuntu.com
    do
      sleep 1
    done
    while ! lxc exec preloaded-${OS} -- apt update
    do
      sleep 1
    done
    lxc exec preloaded-${OS} -- apt dist-upgrade -y
    lxc exec preloaded-${OS} -- apt install -y x11-apps mesa-utils libgl1-mesa-glx
    lxc_stop preloaded-${OS}
    sleep 5   # let it finish stopping
    lxc publish --public preloaded-${OS} --alias gpu-${OS}
    lxc delete preloaded-${OS}
  done
}

# Download a preloaded container and customize it for this machine.
do_localize() {
  case "$1" in
  "") LXDHOST=local;;
  *) LXDHOST=$1;;
  esac

  if ! lxc remote list | grep "$LXDHOST"
  then
    echo "Getting access to remote lxd repository $LXDHOST"
    lxc remote add "$LXDHOST" "$LXDHOST"
  fi

  HOST=$(hostname)
  UID=$(id -u)
  GID=$(id -g)

  case "$DISPLAY" in
  "") echo "Please set DISPLAY first."; exit 1;;
  esac
  XNUM=${DISPLAY#:}

  XSOCK=/tmp/.X11-unix/X$XNUM
  if ! test -S "$XSOCK"
  then
    echo "Please start X first... or something.  Can't open $XSOCK"; exit 1
  fi

  is_nvidia=false
  if lspci | grep VGA | grep -i nvidia
  then
    is_nvidia=true
  fi

  # Create the profile opengl-demo suitable for accessing the GPU
  lxc profile create opengl-demo 2> /dev/null || true
  cat > opengl-demo-profile.tmp <<_EOF_
description: GPU LXD profile for host $HOST
devices:
  gpu:
    type: gpu
  X$XNUM:
    path: $XSOCK
    source: $XSOCK
    type: disk
config:
  environment.DISPLAY: :$XNUM
  raw.idmap: "uid $UID 1000\ngid $GID 1000"
  nvidia.runtime: $is_nvidia
  nvidia.driver.capabilities: "graphics"
_EOF_
  lxc profile edit opengl-demo < opengl-demo-profile.tmp

  lxc_stop opengl-demo-$_os || true
  lxc delete opengl-demo-$_os || true
  lxc launch --profile default --profile opengl-demo "${LXDHOST}":gpu-$_os opengl-demo-$_os

  # If this is an nvidia system, verify generic nvidia driver is working
  if $is_nvidia
  then
    lxc exec opengl-demo-$_os -- "nvidia-smi"
  fi
  # Verify opengl app can run
  lxc exec opengl-demo-$_os -- glxinfo | grep str
  lxc_stop opengl-demo-$_os
}

# Run a command in the localized container.
# By default, just runs glxinfo.
# If 1st arg is -e, container is ephemeral, i.e. any changes are discarded.
do_run() {
  # Clean up after any previous ephemeral run that was interrupted
  lxc delete --force quick 2> /dev/null || true

  local name=quick
  case "$1" in
  --ephemeral|-e) shift; time lxc copy opengl-demo-$_os quick --ephemeral;;
  *) name=opengl-demo-$_os; lxc_stop $name || true; sleep 1;;
  esac

  lxc start $name
  cmd="$*"
  if test "$cmd" = ""
  then
    cmd="glxinfo | grep str"
  fi
  status=0
  # Running xset here is a bit lame, but it's hard to tell if the demo works if screen is blanked
  if ! lxc exec $name -- su ubuntu -c "xset dpms force on && xset -dpms && xset s off && $cmd"
  then
    status=$?
  fi
  lxc_stop $name || true
  return $status
}

#-------------------------------------------------------------------

# Multimachine demo part 1
# Arguments are operating system years (16, 18, 19) to make images for
# Uses the above building blocks to:
# - create an image locally for each supported flavor of ubuntu
# - publish them
do_serve() {
  OSES="$*"
  if test "$OSES" = ""
  then
    echo "Which Ubuntu x.04 do you want to serve?  Give space separate list with 16, 18, and/or 19."
    exit 1
  fi
  # Set up the lxdhost that all machines will pull images from
  do_install
  do_init
  # Tell it to listen on the network (please edit!)
  lxc config set core.https_address [::]:8443
  lxc config set core.trust_password SOME-PASSWORD
  # Populate it with images for the remote machines' operating systems
  do_image $OSES
}

# Multimachine demo part 2
# Run install, init, and localize steps on the given remote hosts.
# FIXME: assumes DISPLAY should be :0 on remote machine...
# and probably assumes you've run xhost + on all the remote machines.
do_multi_setup() {
  for host
  do
    echo "======== Setting up ${host}... ========"
    scp lxc-opengl-demo.sh "${host}":
    ssh -t "${host}" "sh lxc-opengl-demo.sh install && sh lxc-opengl-demo.sh init && DISPLAY=:0 sh lxc-opengl-demo.sh localize $(hostname)"
  done
}

# Multimachine demo part 3
# Run a command in the appropriate image on each given machine concurrently
# Edit to taste (especially cmd).
do_multi_run() {
  local cmd=glxgears
  for host
  do
    echo "======== Running $cmd on ${host}... ========"
    # You might need to do xhost + on the target systems first.
    scp lxc-opengl-demo.sh "${host}":
    ssh -t "${host}" sh -x lxc-opengl-demo.sh run "$cmd" &
  done
  wait
}

# Multimachine demo part 4
# Purge and uninstall lxc/lxd on the given machines
# Optional and traumatic.
do_multi_teardown() {
  for host
  do
    echo "======== Purging ${host}... ========"
    scp lxc-opengl-demo.sh "${host}":
    ssh -t "${host}" 'sh -x lxc-opengl-demo.sh purge && sh -x lxc-opengl-demo.sh uninstall'
  done
}

verb=$1
case "$verb" in
""|help)   usage; exit 0;;
localise)  verb=localize;;
esac

shift
case $verb in
# Simple verbs:
install)    do_install;;
init)       do_init;;
image)      do_image "$@";;
localize)   do_localize "$@";;
run)        do_run "$@";;
uninstall)  do_uninstall;;
purge)      do_purge;;
# Multi-machine demo:
serve)          do_serve "$@";;
multi-setup)    do_multi_setup "$@";;
multi-run)      do_multi_run "$@";;
multi-teardown) do_multi_teardown "$@";;
*)          usage; exit 1;;
esac