Update dea_next's drain script

Change-Id: I26350bb25abce352beb19e34dcfb386f2a85b810

jobs/dea_next/templates/drain.rb

@@ -4,27 +4,47 @@
 
 job_change, hash_change, *updated_packages = ARGV
 
-dea_is_only_changed_package = updated_packages.size == 1 && updated_packages[0] == "dea"
-job_and_hash_unchanged = job_change == "job_unchanged" && hash_change == "hash_unchanged"
+dea_only       = (updated_packages == ["dea_next"])
+warden_only    = (updated_packages == ["warden"])
+dea_and_warden = (updated_packages.sort == ["dea_next", "warden"])
+
+# Must evacuate if job changes, stemcell changes, or a package other than
+# the dea or warden changes
+need_evacuation = (job_change  != "job_unchanged")  ||
+                  (hash_change != "hash_unchanged") ||
+                  !(dea_only || warden_only || dea_and_warden)
+
+dea_pidfile = "/var/vcap/sys/run/dea_next/dea_next.pid"
+warden_pidfile = "/var/vcap/sys/run/warden/warden.pid"
 
-pidfile = "/var/vcap/sys/run/dea/dea.pid"
 default_timeout = 33
 
-if !File.exists?(pidfile)
+if !File.exists?(dea_pidfile)
   puts 0
   exit 0
 end
 
 begin
-  pid = File.read(pidfile).to_i
+  dea_pid = File.read(dea_pidfile).to_i
+  warden_pid = File.read(warden_pidfile).to_i
 
-  if dea_is_only_changed_package && job_and_hash_unchanged
-    Process.kill("KILL", pid)
-    FileUtils.rm_rf(pidfile)
-    puts 0
-  else
-    Process.kill("USR2", pid)
+  if need_evacuation
+    Process.kill("USR2", dea_pid)
     puts (ENV["DEA_DRAIN_TIMEOUT"] || default_timeout).to_i
+    # XXX: The warden should be rolled after the DEA has exited. Unsure how to
+    # make that happen.
+  else
+    Process.kill("KILL", dea_pid)
+    FileUtils.rm_f(dea_pidfile)
+
+    sleep 0.5
+
+    # Persist container state so the DEA can pick up the containers
+    Process.kill("USR2", warden_pid)
+    FileUtils.rm_f(warden_pidfile)
+
+    # Give the warden a bit of time to write out its state
+    puts 1
   end
 
 rescue Errno::ESRCH