#!/usr/bin/env bash

# This script runs the balloon framework on a cluster using the cgroups
# isolator and checks that the framework returns a status of 1.

source $(MESOS_SOURCE_DIR)/support/colors.sh
source $(MESOS_SOURCE_DIR)/support/atexit.sh
source $(MESOS_HELPER_DIR)/colors.sh
source $(MESOS_HELPER_DIR)/atexit.sh

global EXISTING_MEMORY_HIERARCHY := $[cat /proc/mounts | grep memory | cut -f 2 -d ' ]
if [[ -n ${EXISTING_MEMORY_HIERARCHY} ]] {
  # Strip off the subsystem component.
  global TEST_CGROUP_HIERARCHY := $(EXISTING_MEMORY_HIERARCHY%/*)
} else {
  global TEST_CGROUP_HIERARCHY := '/tmp/mesos_test_cgroup'
}
global TEST_CGROUP_ROOT := 'mesos_test'

# Check if the hierarchy exists. If it doesn't, we want to make sure we
# remove it, since we will create it.
global unmount := 'false'
if [[ ! -d ${TEST_CGROUP_HIERARCHY} ]] {
  global unmount := 'true'
}

global MASTER_PID := ''
global AGENT_PID := ''
global MESOS_WORK_DIR := ''
global MESOS_RUNTIME_DIR := ''

# This function ensures that we first kill the agent (if present) and
# then cleanup the cgroups. This is necessary because a running agent
# holds an advisory lock that disallows cleaning up cgroups.
# This function is not pure, but depends on state from the environment
# (e.g. ${AGENT_PID}) because we do not all possible values about when we
# register with 'atexit'.
proc cleanup {
  # Make sure we kill the master on exit.
  if [[ ! -z ${MASTER_PID} ]] {
    kill $(MASTER_PID)
  }

  # Make sure we kill the agent on exit.
  if [[ ! -z ${AGENT_PID} ]] {
    kill $(AGENT_PID)
  }

  # Make sure we cleanup any cgroups we created on exiting.
  find $(TEST_CGROUP_HIERARCHY)/*/$(TEST_CGROUP_ROOT) -mindepth 1 -depth -type d -exec rmdir '{}' '+'

  # Make sure we cleanup the hierarchy, if we created it.
  # NOTE: We do a sleep here, to ensure the hierarchy is not busy.
  if $(unmount) {
   rmdir $(TEST_CGROUP_HIERARCHY)/$(TEST_CGROUP_ROOT) && sleep 1 && umount $(TEST_CGROUP_HIERARCHY) && rmdir $(TEST_CGROUP_HIERARCHY)
  }

  if [[ -d "${MESOS_WORK_DIR}" ]] {
    rm -rf $(MESOS_WORK_DIR);
  }

  if [[ -d "${MESOS_RUNTIME_DIR}" ]] {
    rm -rf $(MESOS_RUNTIME_DIR);
  }
}

atexit cleanup

export LD_LIBRARY_PATH=$(MESOS_BUILD_DIR)/src/.libs
global MASTER := "$(MESOS_SBIN_DIR)/mesos-master"
global AGENT := "$(MESOS_SBIN_DIR)/mesos-agent"
global BALLOON_FRAMEWORK := "$(MESOS_HELPER_DIR)/balloon-framework"

# The mesos binaries expect MESOS_ prefixed environment variables
# to correspond to flags, so we unset these here.
unset MESOS_BUILD_DIR
unset MESOS_SOURCE_DIR
unset MESOS_HELPER_DIR
#unset MESOS_LAUNCHER_DIR # leave this so we can find mesos-fetcher.
unset MESOS_VERBOSE

global MESOS_WORK_DIR := $[mktemp -d -t mesos-XXXXXX]
global MESOS_RUNTIME_DIR := $[mktemp -d -t mesos-XXXXXX]

# Launch master.
$(MASTER) \
    --ip=127.0.0.1 \
    --port=5432 \
    --work_dir=$(MESOS_WORK_DIR) &
global MASTER_PID := $(!)
echo "$(GREEN)Launched master at $(MASTER_PID)$(NORMAL)"
sleep 2

# Check the master is still running after 2 seconds.
kill -0 $(MASTER_PID) >/dev/null !2 > !1
global STATUS := ${?$(?)
if [[ ${STATUS} -ne 0 ]] {
  echo "$(RED)Master crashed; failing test$(NORMAL)"
  exit 2
}

global EXECUTOR_ENVIRONMENT_VARIABLES := ""{\"LD_LIBRARY_PATH\":\"$(LD_LIBRARY_PATH)\"}""

# Launch agent.
$(AGENT) \
    --work_dir=$(MESOS_WORK_DIR) \
    --runtime_dir=$(MESOS_RUNTIME_DIR) \
    --master=127.0.0.1:5432 \
    --isolation=cgroups/mem \
    --cgroups_hierarchy=$(TEST_CGROUP_HIERARCHY) \
    --cgroups_root=$(TEST_CGROUP_ROOT) \
    --executor_environment_variables=$(EXECUTOR_ENVIRONMENT_VARIABLES) \
    --resources="cpus:1;mem:96" &
global AGENT_PID := $(!)
echo "$(GREEN)Launched agent at $(AGENT_PID)$(NORMAL)"
sleep 2

# Check the agent is still running after 2 seconds.
kill -0 $(AGENT_PID) >/dev/null !2 > !1
global STATUS := ${?$(?)
if [[ ${STATUS} -ne 0 ]] {
  echo "$(RED)Slave crashed; failing test$(NORMAL)"
  exit 2
}

# The main event!
$(BALLOON_FRAMEWORK) \
    --master=127.0.0.1:5432 \
    --task_memory_usage_limit=1024MB \
    --task_memory=32MB
global STATUS := ${?$(?)

# Make sure the balloon framework "failed".
if [[ ! ${STATUS} -eq 1 ]] {
  echo "$(RED) Balloon framework returned $(STATUS) not 1$(NORMAL)"
  exit 1
}

# And make sure the agent is still running!
kill -0 $(AGENT_PID) >/dev/null !2 > !1
global STATUS := ${?$(?)
if [[ ${STATUS} -ne 0 ]] {
  echo "$(RED)Slave crashed; failing test$(NORMAL)"
  exit 2
}

exit 0
(CommandList
  children: [
    (C {(source)} {(${ VSub_Name MESOS_SOURCE_DIR) (/support/colors.sh)})
    (C {(source)} {(${ VSub_Name MESOS_SOURCE_DIR) (/support/atexit.sh)})
    (C {(source)} {(${ VSub_Name MESOS_HELPER_DIR) (/colors.sh)})
    (C {(source)} {(${ VSub_Name MESOS_HELPER_DIR) (/atexit.sh)})
    (Assignment
      keyword: Assign_None
      pairs: [
        (assign_pair
          lhs: (LhsName name:EXISTING_MEMORY_HIERARCHY)
          op: Equal
          rhs: 
            {
              (CommandSubPart
                command_list: 
                  (CommandList
                    children: [
                      (Pipeline
                        children: [
                          (C {(cat)} {(/proc/mounts)})
                          (C {(grep)} {(memory)})
                          (C {(cut)} {(-f)} {(2)} {(-d)} {(SQ <" ">)})
                        ]
                        negated: False
                      )
                    ]
                  )
                left_token: <Left_CommandSub "$(">
                spids: [41 65]
              )
            }
          spids: [40]
        )
      ]
      spids: [40]
    )
    (If
      arms: [
        (if_arm
          cond: [
            (Sentence
              child: 
                (DBracket
                  expr: (BoolUnary op_id:BoolUnary_n child:{(${ VSub_Name EXISTING_MEMORY_HIERARCHY)})
                )
              terminator: <Op_Semi ";">
            )
          ]
          action: [
            (Assignment
              keyword: Assign_None
              pairs: [
                (assign_pair
                  lhs: (LhsName name:TEST_CGROUP_HIERARCHY)
                  op: Equal
                  rhs: 
                    {
                      (BracedVarSub
                        token: <VSub_Name EXISTING_MEMORY_HIERARCHY>
                        suffix_op: (StringUnary op_id:VOp1_Percent arg_word:{(Lit_Slash /) ("*")})
                        spids: [88 93]
                      )
                    }
                  spids: [87]
                )
              ]
              spids: [87]
            )
          ]
          spids: [-1 80]
        )
      ]
      else_action: [
        (Assignment
          keyword: Assign_None
          pairs: [
            (assign_pair
              lhs: (LhsName name:TEST_CGROUP_HIERARCHY)
              op: Equal
              rhs: {(/tmp/mesos_test_cgroup)}
              spids: [98]
            )
          ]
          spids: [98]
        )
      ]
      spids: [95 101]
    )
    (Assignment
      keyword: Assign_None
      pairs: [
        (assign_pair
          lhs: (LhsName name:TEST_CGROUP_ROOT)
          op: Equal
          rhs: {(mesos_test)}
          spids: [103]
        )
      ]
      spids: [103]
    )
    (Assignment
      keyword: Assign_None
      pairs: [(assign_pair lhs:(LhsName name:unmount) op:Equal rhs:{(false)} spids:[113])]
      spids: [113]
    )
    (If
      arms: [
        (if_arm
          cond: [
            (Sentence
              child: 
                (DBracket
                  expr: 
                    (LogicalNot
                      child: (BoolUnary op_id:BoolUnary_d child:{(${ VSub_Name TEST_CGROUP_HIERARCHY)})
                    )
                )
              terminator: <Op_Semi ";">
            )
          ]
          action: [
            (Assignment
              keyword: Assign_None
              pairs: [(assign_pair lhs:(LhsName name:unmount) op:Equal rhs:{(true)} spids:[134])]
              spids: [134]
            )
          ]
          spids: [-1 131]
        )
      ]
      spids: [-1 137]
    )
    (Assignment
      keyword: Assign_None
      pairs: [(assign_pair lhs:(LhsName name:MASTER_PID) op:Equal rhs:{(SQ )} spids:[140])]
      spids: [140]
    )
    (Assignment
      keyword: Assign_None
      pairs: [(assign_pair lhs:(LhsName name:AGENT_PID) op:Equal rhs:{(SQ )} spids:[142])]
      spids: [142]
    )
    (Assignment
      keyword: Assign_None
      pairs: [(assign_pair lhs:(LhsName name:MESOS_WORK_DIR) op:Equal rhs:{(SQ )} spids:[144])]
      spids: [144]
    )
    (Assignment
      keyword: Assign_None
      pairs: [(assign_pair lhs:(LhsName name:MESOS_RUNTIME_DIR) op:Equal rhs:{(SQ )} spids:[146])]
      spids: [146]
    )
    (FuncDef
      name: cleanup
      body: 
        (BraceGroup
          children: [
            (If
              arms: [
                (if_arm
                  cond: [
                    (Sentence
                      child: 
                        (DBracket
                          expr: 
                            (LogicalNot
                              child: (BoolUnary op_id:BoolUnary_z child:{(${ VSub_Name MASTER_PID)})
                            )
                        )
                      terminator: <Op_Semi ";">
                    )
                  ]
                  action: [(C {(kill)} {(${ VSub_Name MASTER_PID)})]
                  spids: [-1 195]
                )
              ]
              spids: [-1 205]
            )
            (If
              arms: [
                (if_arm
                  cond: [
                    (Sentence
                      child: 
                        (DBracket
                          expr: 
                            (LogicalNot
                              child: (BoolUnary op_id:BoolUnary_z child:{(${ VSub_Name AGENT_PID)})
                            )
                        )
                      terminator: <Op_Semi ";">
                    )
                  ]
                  action: [(C {(kill)} {(${ VSub_Name AGENT_PID)})]
                  spids: [-1 228]
                )
              ]
              spids: [-1 238]
            )
            (C {(find)} 
              {(${ VSub_Name TEST_CGROUP_HIERARCHY) (/) (Lit_Other "*") (/) 
                (${ VSub_Name TEST_CGROUP_ROOT)
              } {(-mindepth)} {(1)} {(-depth)} {(-type)} {(d)} {(-exec)} {(rmdir)} {(SQ <"{}">)} 
              {(EscapedLiteralPart token:<Lit_EscapedChar "\\+">)}
            )
            (If
              arms: [
                (if_arm
                  cond: [(Sentence child:(C {(${ VSub_Name unmount)}) terminator:<Op_Semi ";">)]
                  action: [
                    (AndOr
                      children: [
                        (C {(rmdir)} 
                          {(${ VSub_Name TEST_CGROUP_HIERARCHY) (/) (${ VSub_Name TEST_CGROUP_ROOT)}
                        )
                        (AndOr
                          children: [
                            (C {(sleep)} {(1)})
                            (AndOr
                              children: [
                                (C {(umount)} {(${ VSub_Name TEST_CGROUP_HIERARCHY)})
                                (C {(rmdir)} {(${ VSub_Name TEST_CGROUP_HIERARCHY)})
                              ]
                              op_id: Op_DAmp
                            )
                          ]
                          op_id: Op_DAmp
                        )
                      ]
                      op_id: Op_DAmp
                    )
                  ]
                  spids: [-1 295]
                )
              ]
              spids: [-1 331]
            )
            (If
              arms: [
                (if_arm
                  cond: [
                    (Sentence
                      child: 
                        (DBracket
                          expr: (BoolUnary op_id:BoolUnary_d child:{(DQ (${ VSub_Name MESOS_WORK_DIR))})
                        )
                      terminator: <Op_Semi ";">
                    )
                  ]
                  action: [
                    (Sentence
                      child: (C {(rm)} {(-rf)} {(${ VSub_Name MESOS_WORK_DIR)})
                      terminator: <Op_Semi ";">
                    )
                  ]
                  spids: [-1 350]
                )
              ]
              spids: [-1 363]
            )
            (If
              arms: [
                (if_arm
                  cond: [
                    (Sentence
                      child: 
                        (DBracket
                          expr: 
                            (BoolUnary
                              op_id: BoolUnary_d
                              child: {(DQ (${ VSub_Name MESOS_RUNTIME_DIR))}
                            )
                        )
                      terminator: <Op_Semi ";">
                    )
                  ]
                  action: [
                    (Sentence
                      child: (C {(rm)} {(-rf)} {(${ VSub_Name MESOS_RUNTIME_DIR)})
                      terminator: <Op_Semi ";">
                    )
                  ]
                  spids: [-1 382]
                )
              ]
              spids: [-1 395]
            )
          ]
          spids: [173]
        )
      spids: [167 172]
    )
    (C {(atexit)} {(cleanup)})
    (C {(export)} {(Lit_VarLike "LD_LIBRARY_PATH=") (${ VSub_Name MESOS_BUILD_DIR) (/src/.libs)})
    (Assignment
      keyword: Assign_None
      pairs: [
        (assign_pair
          lhs: (LhsName name:MASTER)
          op: Equal
          rhs: {(${ VSub_Name MESOS_SBIN_DIR) (/mesos-master)}
          spids: [413]
        )
      ]
      spids: [413]
    )
    (Assignment
      keyword: Assign_None
      pairs: [
        (assign_pair
          lhs: (LhsName name:AGENT)
          op: Equal
          rhs: {(${ VSub_Name MESOS_SBIN_DIR) (/mesos-agent)}
          spids: [419]
        )
      ]
      spids: [419]
    )
    (Assignment
      keyword: Assign_None
      pairs: [
        (assign_pair
          lhs: (LhsName name:BALLOON_FRAMEWORK)
          op: Equal
          rhs: {(${ VSub_Name MESOS_HELPER_DIR) (/balloon-framework)}
          spids: [425]
        )
      ]
      spids: [425]
    )
    (C {(unset)} {(MESOS_BUILD_DIR)})
    (C {(unset)} {(MESOS_SOURCE_DIR)})
    (C {(unset)} {(MESOS_HELPER_DIR)})
    (C {(unset)} {(MESOS_VERBOSE)})
    (Assignment
      keyword: Assign_None
      pairs: [
        (assign_pair
          lhs: (LhsName name:MESOS_WORK_DIR)
          op: Equal
          rhs: 
            {
              (CommandSubPart
                command_list: (CommandList children:[(C {(mktemp)} {(-d)} {(-t)} {(mesos-XXXXXX)})])
                left_token: <Left_Backtick "`">
                spids: [459 467]
              )
            }
          spids: [458]
        )
      ]
      spids: [458]
    )
    (Assignment
      keyword: Assign_None
      pairs: [
        (assign_pair
          lhs: (LhsName name:MESOS_RUNTIME_DIR)
          op: Equal
          rhs: 
            {
              (CommandSubPart
                command_list: (CommandList children:[(C {(mktemp)} {(-d)} {(-t)} {(mesos-XXXXXX)})])
                left_token: <Left_Backtick "`">
                spids: [470 478]
              )
            }
          spids: [469]
        )
      ]
      spids: [469]
    )
    (Sentence
      child: 
        (C {(${ VSub_Name MASTER)} {(--ip) (Lit_Other "=") (127.0.0.1)} 
          {(--port) (Lit_Other "=") (5432)} {(--work_dir) (Lit_Other "=") (${ VSub_Name MESOS_WORK_DIR)}
        )
      terminator: <Op_Amp "&">
    )
    (Assignment
      keyword: Assign_None
      pairs: [
        (assign_pair
          lhs: (LhsName name:MASTER_PID)
          op: Equal
          rhs: {(${ VSub_Bang "!")}
          spids: [510]
        )
      ]
      spids: [510]
    )
    (C {(echo)} 
      {
        (DQ (${ VSub_Name GREEN) ("Launched master at ") (${ VSub_Name MASTER_PID) (${ VSub_Name NORMAL))
      }
    )
    (C {(sleep)} {(2)})
    (SimpleCommand
      words: [{(kill)} {(-0)} {(${ VSub_Name MASTER_PID)}]
      redirects: [
        (Redir op_id:Redir_Great fd:-1 arg_word:{(/dev/null)} spids:[546])
        (Redir op_id:Redir_GreatAnd fd:2 arg_word:{(1)} spids:[549])
      ]
    )
    (Assignment
      keyword: Assign_None
      pairs: [(assign_pair lhs:(LhsName name:STATUS) op:Equal rhs:{(${ VSub_QMark "?")} spids:[552])]
      spids: [552]
    )
    (If
      arms: [
        (if_arm
          cond: [
            (Sentence
              child: 
                (DBracket
                  expr: (BoolBinary op_id:BoolBinary_ne left:{(${ VSub_Name STATUS)} right:{(0)})
                )
              terminator: <Op_Semi ";">
            )
          ]
          action: [
            (C {(echo)} {(DQ (${ VSub_Name RED) ("Master crashed; failing test") (${ VSub_Name NORMAL))})
            (C {(exit)} {(2)})
          ]
          spids: [-1 572]
        )
      ]
      spids: [-1 592]
    )
    (Assignment
      keyword: Assign_None
      pairs: [
        (assign_pair
          lhs: (LhsName name:EXECUTOR_ENVIRONMENT_VARIABLES)
          op: Equal
          rhs: 
            {
              (DQ ("{") (EscapedLiteralPart token:<Lit_EscapedChar "\\\"">) (LD_LIBRARY_PATH) 
                (EscapedLiteralPart token:<Lit_EscapedChar "\\\"">) (":") (EscapedLiteralPart token:<Lit_EscapedChar "\\\"">) (${ VSub_Name LD_LIBRARY_PATH) 
                (EscapedLiteralPart token:<Lit_EscapedChar "\\\"">) ("}")
              )
            }
          spids: [595]
        )
      ]
      spids: [595]
    )
    (Sentence
      child: 
        (C {(${ VSub_Name AGENT)} {(--work_dir) (Lit_Other "=") (${ VSub_Name MESOS_WORK_DIR)} 
          {(--runtime_dir) (Lit_Other "=") (${ VSub_Name MESOS_RUNTIME_DIR)} {(--master) (Lit_Other "=") (127.0.0.1) (Lit_Other ":") (5432)} 
          {(--isolation) (Lit_Other "=") (cgroups/mem)} {(--cgroups_hierarchy) (Lit_Other "=") (${ VSub_Name TEST_CGROUP_HIERARCHY)} 
          {(--cgroups_root) (Lit_Other "=") (${ VSub_Name TEST_CGROUP_ROOT)} {(--executor_environment_variables) (Lit_Other "=") (${ VSub_Name EXECUTOR_ENVIRONMENT_VARIABLES)} 
          {(--resources) (Lit_Other "=") (DQ ("cpus:1;mem:96"))}
        )
      terminator: <Op_Amp "&">
    )
    (Assignment
      keyword: Assign_None
      pairs: [
        (assign_pair
          lhs: (LhsName name:AGENT_PID)
          op: Equal
          rhs: {(${ VSub_Bang "!")}
          spids: [682]
        )
      ]
      spids: [682]
    )
    (C {(echo)} 
      {
        (DQ (${ VSub_Name GREEN) ("Launched agent at ") (${ VSub_Name AGENT_PID) (${ VSub_Name NORMAL))
      }
    )
    (C {(sleep)} {(2)})
    (SimpleCommand
      words: [{(kill)} {(-0)} {(${ VSub_Name AGENT_PID)}]
      redirects: [
        (Redir op_id:Redir_Great fd:-1 arg_word:{(/dev/null)} spids:[718])
        (Redir op_id:Redir_GreatAnd fd:2 arg_word:{(1)} spids:[721])
      ]
    )
    (Assignment
      keyword: Assign_None
      pairs: [(assign_pair lhs:(LhsName name:STATUS) op:Equal rhs:{(${ VSub_QMark "?")} spids:[724])]
      spids: [724]
    )
    (If
      arms: [
        (if_arm
          cond: [
            (Sentence
              child: 
                (DBracket
                  expr: (BoolBinary op_id:BoolBinary_ne left:{(${ VSub_Name STATUS)} right:{(0)})
                )
              terminator: <Op_Semi ";">
            )
          ]
          action: [
            (C {(echo)} {(DQ (${ VSub_Name RED) ("Slave crashed; failing test") (${ VSub_Name NORMAL))})
            (C {(exit)} {(2)})
          ]
          spids: [-1 744]
        )
      ]
      spids: [-1 764]
    )
    (C {(${ VSub_Name BALLOON_FRAMEWORK)} 
      {(--master) (Lit_Other "=") (127.0.0.1) (Lit_Other ":") (5432)} {(--task_memory_usage_limit) (Lit_Other "=") (1024MB)} {(--task_memory) (Lit_Other "=") (32MB)}
    )
    (Assignment
      keyword: Assign_None
      pairs: [(assign_pair lhs:(LhsName name:STATUS) op:Equal rhs:{(${ VSub_QMark "?")} spids:[794])]
      spids: [794]
    )
    (If
      arms: [
        (if_arm
          cond: [
            (Sentence
              child: 
                (DBracket
                  expr: 
                    (LogicalNot
                      child: (BoolBinary op_id:BoolBinary_eq left:{(${ VSub_Name STATUS)} right:{(1)})
                    )
                )
              terminator: <Op_Semi ";">
            )
          ]
          action: [
            (C {(echo)} 
              {
                (DQ (${ VSub_Name RED) (" Balloon framework returned ") (${ VSub_Name STATUS) (" not 1") 
                  (${ VSub_Name NORMAL)
                )
              }
            )
            (C {(exit)} {(1)})
          ]
          spids: [-1 820]
        )
      ]
      spids: [-1 844]
    )
    (SimpleCommand
      words: [{(kill)} {(-0)} {(${ VSub_Name AGENT_PID)}]
      redirects: [
        (Redir op_id:Redir_Great fd:-1 arg_word:{(/dev/null)} spids:[858])
        (Redir op_id:Redir_GreatAnd fd:2 arg_word:{(1)} spids:[861])
      ]
    )
    (Assignment
      keyword: Assign_None
      pairs: [(assign_pair lhs:(LhsName name:STATUS) op:Equal rhs:{(${ VSub_QMark "?")} spids:[864])]
      spids: [864]
    )
    (If
      arms: [
        (if_arm
          cond: [
            (Sentence
              child: 
                (DBracket
                  expr: (BoolBinary op_id:BoolBinary_ne left:{(${ VSub_Name STATUS)} right:{(0)})
                )
              terminator: <Op_Semi ";">
            )
          ]
          action: [
            (C {(echo)} {(DQ (${ VSub_Name RED) ("Slave crashed; failing test") (${ VSub_Name NORMAL))})
            (C {(exit)} {(2)})
          ]
          spids: [-1 884]
        )
      ]
      spids: [-1 904]
    )
    (C {(exit)} {(0)})
  ]
)