Remove tmp image if the download operation fails

prog/download_boot_image.rb

@@ -26,11 +26,9 @@ def blob_storage_client
   end
 
   label def wait_draining
-    unless vm_host.vms.empty?
-      nap 15
-    end
-    sshable.cmd("sudo rm -f /var/storage/images/#{image_name.shellescape}.raw")
-    hop_download
+    hop_download if vm_host.vms.empty?
+
+    nap 15
   end
 
   label def download
@@ -73,6 +71,9 @@ def blob_storage_client
   end
 
   label def activate_host
+    # Verify that the image was downloaded
+    sshable.cmd("ls /var/storage/images/#{image_name}.raw")
+
     vm_host.update(allocation_state: "accepting")
 
     pop "#{image_name} downloaded"

rhizome/host/bin/download-boot-image

@@ -14,4 +14,4 @@ require_relative "../lib/vm_setup"
 certs = $stdin.read
 ca_path = "/usr/lib/ssl/certs/ubicloud_images_blob_storage_certs.crt"
 safe_write_to_file(ca_path, certs)
-VmSetup.new("").download_boot_image(boot_image, custom_url: custom_url, ca_path: ca_path)
+VmSetup.new("").download_boot_image(boot_image, force: true, custom_url: custom_url, ca_path: ca_path)

rhizome/host/lib/vm_setup.rb

@@ -459,7 +459,7 @@ def storage(storage_params, storage_secrets, prep)
     }
   end
 
-  def download_boot_image(boot_image, custom_url: nil, ca_path: nil)
+  def download_boot_image(boot_image, force: false, custom_url: nil, ca_path: nil)
     urls = {
       "ubuntu-jammy" => "https://cloud-images.ubuntu.com/releases/jammy/release-20231010/ubuntu-22.04-server-cloudimg-#{Arch.render(x64: "amd64")}.img",
       "almalinux-9.1" => Arch.render(x64: "x86_64", arm64: "aarch64").yield_self { "https://repo.almalinux.org/almalinux/9/cloud/#{_1}/images/AlmaLinux-9-GenericCloud-latest.#{_1}.qcow2" },
@@ -470,42 +470,42 @@ def download_boot_image(boot_image, custom_url: nil, ca_path: nil)
 
     download = urls.fetch(boot_image) || custom_url
     image_path = vp.image_path(boot_image)
-    unless File.exist?(image_path)
-      fail "Must provide custom_url for #{boot_image} image" if download.nil?
-      FileUtils.mkdir_p vp.image_root
-
-      # If image URL has query parameter such as SAS token, File.extname returns
-      # it too. We need to remove them and only get extension.
-      image_ext = File.extname(URI.parse(download).path)
-      initial_format = case image_ext
-      when ".qcow2", ".img"
-        "qcow2"
-      when ".vhd"
-        "vpc"
-      when ".raw"
-        "raw"
-      else
-        fail "Unsupported boot_image format: #{image_ext}"
-      end
+    return if File.exist?(image_path) && !force
+
+    fail "Must provide custom_url for #{boot_image} image" if download.nil?
+    FileUtils.mkdir_p vp.image_root
+
+    # If image URL has query parameter such as SAS token, File.extname returns
+    # it too. We need to remove them and only get extension.
+    image_ext = File.extname(URI.parse(download).path)
+    initial_format = case image_ext
+    when ".qcow2", ".img"
+      "qcow2"
+    when ".vhd"
+      "vpc"
+    when ".raw"
+      "raw"
+    else
+      fail "Unsupported boot_image format: #{image_ext}"
+    end
 
-      # Use of File::EXCL provokes a crash rather than a race
-      # condition if two VMs are lazily getting their images at the
-      # same time.
-      temp_path = File.join(vp.image_root, boot_image + image_ext + ".tmp")
+    File.open(File.join(vp.image_root, boot_image + ".lock"), File::RDWR | File::CREAT) do |lock|
+      fail "Another vm is downloading #{boot_image}" unless lock.flock(File::LOCK_EX | File::LOCK_NB)
+
+      download_path = File.join(vp.image_root, boot_image + image_ext + ".tmp")
       ca_arg = ca_path ? " --cacert #{ca_path.shellescape}" : ""
-      File.open(temp_path, File::RDWR | File::CREAT | File::EXCL, 0o644) do
-        r "curl -f -L10 -o #{temp_path.shellescape} #{download.shellescape}#{ca_arg}"
+      File.open(download_path, File::RDWR | File::CREAT, 0o644) do
+        r "curl -f -L10 -o #{download_path.shellescape} #{download.shellescape}#{ca_arg}"
       end
 
-      if initial_format == "raw"
-        File.rename(temp_path, image_path)
-      else
+      temp_path = File.join(vp.image_root, boot_image + ".raw.tmp")
+      if initial_format != "raw"
         # Images are presumed to be atomically renamed into the path,
         # i.e. no partial images will be passed to qemu-image.
-        r "qemu-img convert -p -f #{initial_format.shellescape} -O raw #{temp_path.shellescape} #{image_path.shellescape}"
+        r "qemu-img convert -p -f #{initial_format.shellescape} -O raw #{download_path.shellescape} #{temp_path.shellescape}"
+        rm_if_exists(download_path)
       end
-
-      rm_if_exists(temp_path)
+      File.rename(temp_path, image_path)
     end
   end
 

rhizome/host/spec/vm_setup_spec.rb

@@ -56,44 +56,64 @@ def key_wrapping_secrets
   describe "#download_boot_image" do
     it "can download an image" do
       expect(File).to receive(:exist?).with("/var/storage/images/ubuntu-jammy.raw").and_return(false)
+      expect(File).to receive(:open).with("/var/storage/images/ubuntu-jammy.lock", File::RDWR | File::CREAT).and_yield(instance_double(File, flock: true))
       expect(File).to receive(:open) do |path, *_args|
         expect(path).to eq("/var/storage/images/ubuntu-jammy.img.tmp")
       end.and_yield
       expect(FileUtils).to receive(:mkdir_p).with("/var/storage/images/")
       expect(Arch).to receive(:render).and_return("amd64").at_least(:once)
       expect(vs).to receive(:r).with("curl -f -L10 -o /var/storage/images/ubuntu-jammy.img.tmp https://cloud-images.ubuntu.com/releases/jammy/release-20231010/ubuntu-22.04-server-cloudimg-amd64.img")
-      expect(vs).to receive(:r).with("qemu-img convert -p -f qcow2 -O raw /var/storage/images/ubuntu-jammy.img.tmp /var/storage/images/ubuntu-jammy.raw")
+      expect(vs).to receive(:r).with("qemu-img convert -p -f qcow2 -O raw /var/storage/images/ubuntu-jammy.img.tmp /var/storage/images/ubuntu-jammy.raw.tmp")
       expect(FileUtils).to receive(:rm_r).with("/var/storage/images/ubuntu-jammy.img.tmp")
+      expect(File).to receive(:rename).with("/var/storage/images/ubuntu-jammy.raw.tmp", "/var/storage/images/ubuntu-jammy.raw")
 
       vs.download_boot_image("ubuntu-jammy")
     end
 
     it "can download vhd image with custom URL that has query params using curl" do
       expect(File).to receive(:exist?).with("/var/storage/images/github-ubuntu-2204.raw").and_return(false)
+      expect(File).to receive(:open).with("/var/storage/images/github-ubuntu-2204.lock", File::RDWR | File::CREAT).and_yield(instance_double(File, flock: true))
       expect(File).to receive(:open) do |path, *_args|
         expect(path).to eq("/var/storage/images/github-ubuntu-2204.vhd.tmp")
       end.and_yield
       expect(FileUtils).to receive(:mkdir_p).with("/var/storage/images/")
       expect(vs).to receive(:r).with("curl -f -L10 -o /var/storage/images/github-ubuntu-2204.vhd.tmp http://minio.ubicloud.com:9000/ubicloud-images/ubuntu-22.04-x64.vhd\\?X-Amz-Algorithm\\=AWS4-HMAC-SHA256\\&X-Amz-Credential\\=user\\%2F20240112\\%2Fus-east-1\\%2Fs3\\%2Faws4_request\\&X-Amz-Date\\=20240112T132931Z\\&X-Amz-Expires\\=3600\\&X-Amz-SignedHeaders\\=host\\&X-Amz-Signature\\=aabbcc")
-      expect(vs).to receive(:r).with("qemu-img convert -p -f vpc -O raw /var/storage/images/github-ubuntu-2204.vhd.tmp /var/storage/images/github-ubuntu-2204.raw")
+      expect(vs).to receive(:r).with("qemu-img convert -p -f vpc -O raw /var/storage/images/github-ubuntu-2204.vhd.tmp /var/storage/images/github-ubuntu-2204.raw.tmp")
       expect(FileUtils).to receive(:rm_r).with("/var/storage/images/github-ubuntu-2204.vhd.tmp")
+      expect(File).to receive(:rename).with("/var/storage/images/github-ubuntu-2204.raw.tmp", "/var/storage/images/github-ubuntu-2204.raw")
 
       vs.download_boot_image("github-ubuntu-2204", custom_url: "http://minio.ubicloud.com:9000/ubicloud-images/ubuntu-22.04-x64.vhd?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=user%2F20240112%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240112T132931Z&X-Amz-Expires=3600&X-Amz-SignedHeaders=host&X-Amz-Signature=aabbcc")
     end
 
     it "does not convert image if it's in raw format already" do
       expect(File).to receive(:exist?).with("/var/storage/images/github-ubuntu-2204.raw").and_return(false)
+      expect(File).to receive(:open).with("/var/storage/images/github-ubuntu-2204.lock", File::RDWR | File::CREAT).and_yield(instance_double(File, flock: true))
       expect(File).to receive(:open) do |path, *_args|
         expect(path).to eq("/var/storage/images/github-ubuntu-2204.raw.tmp")
       end.and_yield
       expect(FileUtils).to receive(:mkdir_p).with("/var/storage/images/")
       expect(vs).to receive(:r).with("curl -f -L10 -o /var/storage/images/github-ubuntu-2204.raw.tmp http://minio.ubicloud.com:9000/ubicloud-images/ubuntu-22.04-x64.raw\\?X-Amz-Algorithm\\=AWS4-HMAC-SHA256\\&X-Amz-Credential\\=user\\%2F20240112\\%2Fus-east-1\\%2Fs3\\%2Faws4_request\\&X-Amz-Date\\=20240112T132931Z\\&X-Amz-Expires\\=3600\\&X-Amz-SignedHeaders\\=host\\&X-Amz-Signature\\=aabbcc")
       expect(File).to receive(:rename).with("/var/storage/images/github-ubuntu-2204.raw.tmp", "/var/storage/images/github-ubuntu-2204.raw")
-      expect(FileUtils).to receive(:rm_r).with("/var/storage/images/github-ubuntu-2204.raw.tmp")
 
       vs.download_boot_image("github-ubuntu-2204", custom_url: "http://minio.ubicloud.com:9000/ubicloud-images/ubuntu-22.04-x64.raw?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=user%2F20240112%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240112T132931Z&X-Amz-Expires=3600&X-Amz-SignedHeaders=host&X-Amz-Signature=aabbcc")
     end
 
+    it "can download the image with force even if it exists" do
+      expect(File).to receive(:exist?).with("/var/storage/images/ubuntu-jammy.raw").and_return(true)
+      expect(File).to receive(:open).with("/var/storage/images/ubuntu-jammy.lock", File::RDWR | File::CREAT).and_yield(instance_double(File, flock: true))
+      expect(File).to receive(:open) do |path, *_args|
+        expect(path).to eq("/var/storage/images/ubuntu-jammy.img.tmp")
+      end.and_yield
+      expect(FileUtils).to receive(:mkdir_p).with("/var/storage/images/")
+      expect(Arch).to receive(:render).and_return("amd64").at_least(:once)
+      expect(vs).to receive(:r).with("curl -f -L10 -o /var/storage/images/ubuntu-jammy.img.tmp https://cloud-images.ubuntu.com/releases/jammy/release-20231010/ubuntu-22.04-server-cloudimg-amd64.img")
+      expect(vs).to receive(:r).with("qemu-img convert -p -f qcow2 -O raw /var/storage/images/ubuntu-jammy.img.tmp /var/storage/images/ubuntu-jammy.raw.tmp")
+      expect(FileUtils).to receive(:rm_r).with("/var/storage/images/ubuntu-jammy.img.tmp")
+      expect(File).to receive(:rename).with("/var/storage/images/ubuntu-jammy.raw.tmp", "/var/storage/images/ubuntu-jammy.raw")
+
+      vs.download_boot_image("ubuntu-jammy", force: true)
+    end
+
     it "can use an image that's already downloaded" do
       expect(File).to receive(:exist?).with("/var/storage/images/almalinux-9.1.raw").and_return(true)
       vs.download_boot_image("almalinux-9.1")
@@ -109,6 +129,15 @@ def key_wrapping_secrets
       expect(FileUtils).to receive(:mkdir_p).with("/var/storage/images/")
       expect { vs.download_boot_image("github-ubuntu-2204", custom_url: "https://example.com/ubuntu.iso") }.to raise_error RuntimeError, "Unsupported boot_image format: .iso"
     end
+
+    it "fails if another vm is already downloading the image" do
+      expect(File).to receive(:exist?).with("/var/storage/images/ubuntu-jammy.raw").and_return(false)
+      expect(FileUtils).to receive(:mkdir_p).with("/var/storage/images/")
+      expect(Arch).to receive(:render).and_return("amd64").at_least(:once)
+      expect(File).to receive(:open).with("/var/storage/images/ubuntu-jammy.lock", File::RDWR | File::CREAT).and_yield(instance_double(File, flock: false))
+
+      expect { vs.download_boot_image("ubuntu-jammy") }.to raise_error RuntimeError, "Another vm is downloading ubuntu-jammy"
+    end
   end
 
   describe "#purge_storage" do

spec/prog/download_boot_image_spec.rb

@@ -28,7 +28,6 @@
 
     it "hops if it's drained" do
       expect(vm_host).to receive(:vms).and_return([])
-      expect(sshable).to receive(:cmd).with("sudo rm -f /var/storage/images/my-image.raw")
       expect { dbi.wait_draining }.to hop("download")
     end
   end
@@ -99,6 +98,7 @@
 
   describe "#activate_host" do
     it "activates vm host again" do
+      expect(sshable).to receive(:cmd).with("ls /var/storage/images/my-image.raw")
       expect {
         expect { dbi.activate_host }.to exit({"msg" => "my-image downloaded"})
       }.to change { vm_host.reload.allocation_state }.from("unprepared").to("accepting")