Add optional per-job disk_size_gb to provision_service task

README.md

@@ -291,6 +291,10 @@ To request a specific VM machine type for a job (for example, a larger VM for me
 
 The service has to support the field first. A service that does not yet know about `machine_type` ignores it without any error, so the VM silently stays at the service default size. A service that does support it returns HTTP 400 for a value it rejects, but the facade currently reports some of those rejections as a 5xx, which this task retries with backoff for a few minutes before failing.
 
+To request a specific boot disk size for a job (for example, when a test suite runs out of disk space on the default-sized VM), pass the optional `disk_size_gb` task parameter or set the `PROVISION_DISK_SIZE_GB` environment variable (the parameter wins; a blank value counts as unset). With Litmus, which cannot pass extra task parameters, put `provision_disk_size_gb: <size>` under `params` in `provision.yaml` when using `litmus:provision_list` (Litmus exports each `params` key as an upper-cased environment variable) or set `PROVISION_DISK_SIZE_GB` in the job's environment. It is sent to the service as `disk_size_gb` on each VM and is omitted when unset, so the service keeps its default (the source image's own size).
+
+The service has to support the field first. A service that does not yet know about `disk_size_gb` ignores it without any error, so the VM silently stays at its default disk size. A service that does support it returns HTTP 400 for a value it rejects, but the facade currently reports some of those rejections as a 5xx, which this task retries with backoff for a few minutes before failing.
+
 Manual invocation of the provision service task from a workflow can be done using:
 
 ```ruby

spec/tasks/provision_service_spec.rb

@@ -9,10 +9,14 @@
 
 describe 'ProvisionService' do
   around(:each) do |example|
-    saved = ENV.delete('PROVISION_MACHINE_TYPE')
+    saved = ENV.slice('PROVISION_MACHINE_TYPE', 'PROVISION_DISK_SIZE_GB')
+    ENV.delete('PROVISION_MACHINE_TYPE')
+    ENV.delete('PROVISION_DISK_SIZE_GB')
     example.run
   ensure
-    saved ? ENV['PROVISION_MACHINE_TYPE'] = saved : ENV.delete('PROVISION_MACHINE_TYPE')
+    ENV.delete('PROVISION_MACHINE_TYPE')
+    ENV.delete('PROVISION_DISK_SIZE_GB')
+    saved.each { |key, value| ENV[key] = value }
   end
 
   describe '.run' do
@@ -74,15 +78,23 @@
       allow($stdin).to receive(:read).and_return('{"action":"provision","platform":"centos-8"}')
       runner = ProvisionService.new
       allow(ProvisionService).to receive(:new).and_return(runner)
-      expect(runner).to receive(:provision).with('centos-8', instance_of(InventoryHelper), nil, nil, nil).and_return({ status: 'ok', node_name: 'centos-8' })
+      expect(runner).to receive(:provision).with('centos-8', instance_of(InventoryHelper), nil, nil, nil, nil).and_return({ status: 'ok', node_name: 'centos-8' })
       expect { ProvisionService.run }.to raise_error(SystemExit) { |e| expect(e.status).to eq(0) }
     end
 
     it 'passes machine_type through to provision' do
       allow($stdin).to receive(:read).and_return('{"action":"provision","platform":"centos-8","machine_type":"n1-standard-4"}')
       runner = ProvisionService.new
       allow(ProvisionService).to receive(:new).and_return(runner)
-      expect(runner).to receive(:provision).with('centos-8', instance_of(InventoryHelper), nil, nil, 'n1-standard-4').and_return({ status: 'ok', node_name: 'centos-8' })
+      expect(runner).to receive(:provision).with('centos-8', instance_of(InventoryHelper), nil, nil, 'n1-standard-4', nil).and_return({ status: 'ok', node_name: 'centos-8' })
+      expect { ProvisionService.run }.to raise_error(SystemExit) { |e| expect(e.status).to eq(0) }
+    end
+
+    it 'passes disk_size_gb through to provision' do
+      allow($stdin).to receive(:read).and_return('{"action":"provision","platform":"centos-8","disk_size_gb":100}')
+      runner = ProvisionService.new
+      allow(ProvisionService).to receive(:new).and_return(runner)
+      expect(runner).to receive(:provision).with('centos-8', instance_of(InventoryHelper), nil, nil, nil, 100).and_return({ status: 'ok', node_name: 'centos-8' })
       expect { ProvisionService.run }.to raise_error(SystemExit) { |e| expect(e.status).to eq(0) }
     end
 
@@ -251,6 +263,45 @@
       end
     end
 
+    context 'when a disk_size_gb is requested' do
+      let(:request_body) do
+        '{"url":"https://api.github.com/repos/puppetlabs/puppetlabs-iis/actions/runs/1234567890","VMs":[{"cloud":null,"region":null,"zone":null,"images":["centos-8"],"disk_size_gb":100}]}'
+      end
+
+      it 'sends it to the service when passed as a parameter' do
+        stub = stub_request(:post, service_url).with(body: request_body).to_return(status: 200, body: response_body.to_json, headers: {})
+        allow(File).to receive(:open)
+        expect(provision_service.provision(platform, inventory, vars, retry_attempts, nil, 100)[:status]).to eq('ok')
+        expect(stub).to have_been_requested
+      end
+
+      it 'falls back to the PROVISION_DISK_SIZE_GB environment variable' do
+        stub = stub_request(:post, service_url).with(body: request_body).to_return(status: 200, body: response_body.to_json, headers: {})
+        allow(File).to receive(:open)
+        ENV['PROVISION_DISK_SIZE_GB'] = '100'
+        expect(provision_service.provision(platform, inventory, vars, retry_attempts)[:status]).to eq('ok')
+        expect(stub).to have_been_requested
+      end
+
+      it 'prefers the parameter over the PROVISION_DISK_SIZE_GB environment variable' do
+        stub = stub_request(:post, service_url).with(body: request_body).to_return(status: 200, body: response_body.to_json, headers: {})
+        allow(File).to receive(:open)
+        ENV['PROVISION_DISK_SIZE_GB'] = '50'
+        expect(provision_service.provision(platform, inventory, vars, retry_attempts, nil, 100)[:status]).to eq('ok')
+        expect(stub).to have_been_requested
+      end
+
+      it 'sends nothing for a blank parameter and blank environment variable' do
+        stub = stub_request(:post, service_url)
+               .with { |request| JSON.parse(request.body)['VMs'].none? { |vm| vm.key?('disk_size_gb') } }
+               .to_return(status: 200, body: response_body.to_json, headers: {})
+        allow(File).to receive(:open)
+        ENV['PROVISION_DISK_SIZE_GB'] = ' '
+        expect(provision_service.provision(platform, inventory, vars, retry_attempts, nil, '  ')[:status]).to eq('ok')
+        expect(stub).to have_been_requested
+      end
+    end
+
     context 'when the service returns a 5xx and then succeeds' do
       it 'retries with a backoff and returns a valid response' do
         stub_request(:post, service_url)
@@ -331,6 +382,18 @@
       expect(svc.platform_to_cloud_request_parameters({ images: 'centos-8' }, nil, nil, nil, 'n1-standard-4')[:machine_type]).to eq('n1-standard-4')
     end
 
+    it 'omits disk_size_gb unless one is given' do
+      expect(svc.platform_to_cloud_request_parameters('centos-8', nil, nil, nil)).not_to have_key(:disk_size_gb)
+      expect(svc.platform_to_cloud_request_parameters('centos-8', nil, nil, nil, nil, '')).not_to have_key(:disk_size_gb)
+      expect(svc.platform_to_cloud_request_parameters('centos-8', nil, nil, nil, nil, ' ')).not_to have_key(:disk_size_gb)
+    end
+
+    it 'adds disk_size_gb as an Integer for String, Array and Hash platforms' do
+      expect(svc.platform_to_cloud_request_parameters('centos-8', nil, nil, nil, nil, 100)[:disk_size_gb]).to eq(100)
+      expect(svc.platform_to_cloud_request_parameters(['centos-8'], nil, nil, nil, nil, '100')[:disk_size_gb]).to eq(100)
+      expect(svc.platform_to_cloud_request_parameters({ images: 'centos-8' }, nil, nil, nil, nil, 100)[:disk_size_gb]).to eq(100)
+    end
+
     it 'handles a Hash platform and overrides the cloud' do
       platform = { images: ['centos-8'] }
       result = svc.platform_to_cloud_request_parameters(platform, 'aws', nil, nil)

tasks/provision_service.json

@@ -28,6 +28,10 @@
       "description": "Machine type to request for the VM (e.g. n1-standard-4). Defaults to the PROVISION_MACHINE_TYPE environment variable, then to the service's own default",
       "type": "Optional[String[1]]"
     },
+    "disk_size_gb": {
+      "description": "Boot disk size in GB to request for the VM. Defaults to the PROVISION_DISK_SIZE_GB environment variable, then to the service's own default (the source image's own size)",
+      "type": "Optional[Integer[1]]"
+    },
     "retry_attempts": {
       "description": "The number of times to retry the provisioning if it fails",
       "type": "Optional[Integer[1]]",

tasks/provision_service.rb

@@ -37,7 +37,7 @@ def default_uri
     'https://facade-release-6f3kfepqcq-ew.a.run.app/v1/provision'
   end
 
-  def platform_to_cloud_request_parameters(platform, cloud, region, zone, machine_type = nil)
+  def platform_to_cloud_request_parameters(platform, cloud, region, zone, machine_type = nil, disk_size_gb = nil)
     params = case platform
              when String
                { cloud:, region:, zone:, images: [platform] }
@@ -51,6 +51,8 @@ def platform_to_cloud_request_parameters(platform, cloud, region, zone, machine_
     # Only sent when requested, so the request is unchanged for callers that don't set it.
     machine_type = machine_type.to_s.strip
     params[:machine_type] = machine_type unless machine_type.empty?
+    disk_size_gb = disk_size_gb.to_s.strip
+    params[:disk_size_gb] = disk_size_gb.to_i unless disk_size_gb.empty?
     params
   end
 
@@ -113,7 +115,7 @@ def retry_delay(attempt)
     delay + rand(0..(delay / 2))
   end
 
-  def provision(platform, inventory, vars, retry_attempts, machine_type = nil)
+  def provision(platform, inventory, vars, retry_attempts, machine_type = nil, disk_size_gb = nil)
     # Call the provision service with the information necessary and write the inventory file locally
 
     if ENV['GITHUB_RUN_ID']
@@ -126,14 +128,15 @@ def provision(platform, inventory, vars, retry_attempts, machine_type = nil)
     region = ENV.fetch('REGION', nil)
     zone = ENV.fetch('ZONE', nil)
     machine_type = ENV.fetch('PROVISION_MACHINE_TYPE', nil) if machine_type.to_s.strip.empty?
+    disk_size_gb = ENV.fetch('PROVISION_DISK_SIZE_GB', nil) if disk_size_gb.to_s.strip.empty?
     if job_url.nil? && vars
       data = JSON.parse(vars.tr(';', ','))
       job_url = data['job_url']
     end
     attempt = 0
     begin
       attempt += 1
-      params = platform_to_cloud_request_parameters(platform, cloud, region, zone, machine_type)
+      params = platform_to_cloud_request_parameters(platform, cloud, region, zone, machine_type, disk_size_gb)
       response = invoke_cloud_request(params, uri, job_url, 'post', retry_attempts)
       response_hash = YAML.safe_load(response)
       # Knock the response for validity to make sure return payload is expected.
@@ -190,13 +193,16 @@ def tear_down(node_name, inventory, _vars, retry_attempts)
   # - retry_attempts: The number of retry attempts for provisioning or tearing down.
   # - machine_type: Optional VM machine type to request from the service (e.g. 'n1-standard-4'). Falls back to the PROVISION_MACHINE_TYPE
   #   environment variable; when neither is set the service chooses its default.
+  # - disk_size_gb: Optional boot disk size (in GB) to request from the service. Falls back to the PROVISION_DISK_SIZE_GB environment
+  #   variable; when neither is set the service chooses its default (the source image's own size).
   # - inventory_location: The location of the inventory file. If not provided, defaults to './spec/fixtures/litmus_inventory.yaml'.
   # The result of the action is printed to stdout in JSON format.
   # Exits with status 0 on success, or 1 on failure.
   def self.run
     params = JSON.parse($stdin.read)
     params.transform_keys!(&:to_sym)
-    action, node_name, platform, vars, retry_attempts, inventory_location, machine_type = params.values_at(:action, :node_name, :platform, :vars, :retry_attempts, :inventory, :machine_type)
+    action, node_name, platform, vars, retry_attempts, inventory_location, machine_type, disk_size_gb =
+      params.values_at(:action, :node_name, :platform, :vars, :retry_attempts, :inventory, :machine_type, :disk_size_gb)
 
     inventory_location ||= File.join(Dir.pwd, '/spec/fixtures/litmus_inventory.yaml')
     inventory = InventoryHelper.open(inventory_location)
@@ -207,7 +213,7 @@ def self.run
       when 'provision'
         raise 'specify a platform when provisioning' if platform.to_s.empty?
 
-        result = runner.provision(platform, inventory, vars, retry_attempts, machine_type)
+        result = runner.provision(platform, inventory, vars, retry_attempts, machine_type, disk_size_gb)
       when 'tear_down'
         raise 'specify a node_name when tearing down' if node_name.nil?